%outputs_5, %control_6 = tf_executor.island wraps "tf.PreventGradient"(%outputs_2) {device = "", message = "The jax2tf-converted function does not support gradients. Use `with_gradient` parameter to enable gradients"} : (tensor<2x3xi32>) -> tensor<2x3xi32>
      %outputs_7, %control_8 = tf_executor.island wraps "tf.Identity"(%outputs_5) {device = ""} : (tensor<2x3xi32>) -> tensor<2x3xi32>

      %outputs_4, %control_5 = tf_executor.island wraps "tf.BiasAdd"(%outputs_2, %outputs) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<32xf32>) -> tensor<*xf32>
      %outputs_6, %control_7 = tf_executor.island wraps "tf.Identity"(%outputs_4) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
      tf_executor.fetch %outputs_6 : tensor<*xf32>
    }
    func.return %0 : tensor<*xf32>
  }

      %outputs_2, %control_3 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<-1> : tensor<i64>} : () -> tensor<i64>
      %outputs_4, %control_5 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>

    %outputs_6:2, %control_7 = tf_executor.island wraps "tf.While"(%arg0, %arg1) {body = @body, cond = @cond, is_stateless = false, shape_invariant} : (tensor<i32>, tensor<5xf32>) -> (tensor<i32>, tensor<5xf32>) loc("WhileWithOutputShapes")
    tf_executor.fetch %outputs_2#1, %outputs_4#1, %outputs_6#1 : tensor<5xf32>, tensor<5xf32>, tensor<5xf32>
  }

    %outputs_0, %control_1 = tf_executor.island wraps "tf.LookupTableSizeV2"(%outputs) {device = ""} : (tensor<!tf_type.resource>) -> tensor<i64>
    %outputs_2, %control_3 = tf_executor.island wraps "tf.Const"() {value = dense<1.000000e+00> : tensor<f32>} : () -> tensor<f32>
    %outputs_4, %control_5 = tf_executor.island wraps "tf.Identity"(%outputs_0) : (tensor<i64>) -> tensor<i64>

      %outputs_4, %control_5 = tf_executor.island wraps "tf.Identity"(%outputs_2) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>

      %outputs_17, %control_18 = tf_executor.island wraps "tf.Mul"(%outputs_15, %outputs_0) {T = "tfdtype$DT_FLOAT", _tpu_replicate = "cluster", device = "", name = "mul_1"} : (tensor<0xf32>, tensor<f32>) -> tensor<0xf32>
      %outputs_19, %control_20 = tf_executor.island wraps "tf.AddV2"(%outputs_9, %outputs_17) {T = "tfdtype$DT_FLOAT", _tpu_replicate = "cluster", device = "", name = "add"} : (tensor<0xf32>, tensor<0xf32>) -> tensor<0xf32>

      %outputs_26, %control_27 = tf_executor.island wraps "tf.All"(%outputs_24, %outputs_20) {device = "/device:CPU:0", keep_dims = false} : (tensor<*xi1>, tensor<*xi32>) -> tensor<*xi1>

          %outputs_2, %control_3 = tf_executor.island wraps "tf.XlaSpmdFullToShardShape"(%outputs_0) {dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<2x2xf32>) -> tensor<1x2xf32>

      %outputs_8, %control_9 = tf_executor.island wraps "tf.Const"() {device = "", dtype = !tf_type.string, value = dense<""> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>