// See the function doc in the header file.
def GetNhwcReturnTypeFromNchw: NativeCodeCall<
  "GetNhwcReturnTypeFromNchw((*$0.begin()).getDefiningOp())">;

// When given a DenseIntElementsAttr containing I64 elements, this extracts
// one I32IntegerAttr from the given index.
class GetI32At<int index>: NativeCodeCall<
  "$_builder.getI32IntegerAttr(static_cast<int32_t>(*($0.getValues<int64_t>().begin() + " # index #")))">;

          %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>

  // Check bytes come back in the right order
  %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 16909060]> : tensor<4xi32> } : () -> tensor<4xi32>
  func.return %0 : tensor<4xi32>
}

func.func @i64() -> tensor<4xi64> {
  // CHECK-LABEL: @i64
  // CHECK: value = dense<[1, 2, 3, 72623859790382856]> : tensor<4xi64>
  %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 72623859790382856]> : tensor<4xi64> } : () -> tensor<4xi64>

  // Check bytes come back in the right order
  %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 16909060]> : tensor<4xi32> } : () -> tensor<4xi32>
  func.return %0 : tensor<4xi32>
}

func.func @i64() -> tensor<4xi64> {
  // CHECK-LABEL: @i64
  // CHECK: value = dense<[1, 2, 3, 72623859790382856]> : tensor<4xi64>
  %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 72623859790382856]> : tensor<4xi64> } : () -> tensor<4xi64>

			n = n*10 + int(ch)
		}
		if s0[0] == '-' {
			n = -n
		}
		return n, nil
	}

	// Slow path for invalid, big, or underscored integers.
	i64, err := ParseInt(s, 10, 0)
	if nerr, ok := err.(*NumError); ok {
		nerr.Func = fnAtoi
	}
	return int(i64), err
}

// underscoreOK reports whether the underscores in s are allowed.
// Checking them in this one function lets all the parsers skip over them simply.

      %1 = stablehlo.constant dense<[[-0.630731344, 0.54962182], [0.180364341, -0.764542698]]> : tensor<2x2xf32>
      %2 = "tf.XlaCallModule"(%arg0, %1, %0) <{Sout = [#tf_type.shape<?x2>], module = "", version = 9 : i64}> {_entry_function = @composite_fn_1, _original_entry_function = "composite_fn_1", _tfl_quant_trait = "fully_quantizable"} : (tensor<?x2xf32>, tensor<2x2xf32>, tensor<2xf32>) -> tensor<?x2xf32>
      return %2 : tensor<?x2xf32>
    }

    %0 = "tf.XlaCallModule"(%output, %cst) <{Sout = [#tf_type.shape<10x1x3>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_with_relu_fn_1, _original_entry_function = "composite_dot_general_with_relu_fn_1", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait...

        Sout = [#tf_type.shape<?x10>], dim_args_spec = [],
        disabled_checks = [], function_list = [],
        has_token_input_output = false, module = "", platforms = [],
        version = 5 : i64
    }> {
        _entry_function = @composite_dot_general_fn_1,
        _original_entry_function = "composite_dot_general_fn_1",
        _tfl_quant_trait = "fully_quantizable",

// CHECK: return %[[SUB_0]] : tensor<15x28x28x1xf32>
}

// CHECK-LABEL: depth_to_space
func.func @depth_to_space(%arg0: tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32> {
  %0 = "tf.DepthToSpace"(%arg0) {block_size = 2: i64,  data_format = "NHWC"}: (tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32>
  func.return %0 : tensor<1x2x2x1xf32>

    }) : () -> tensor<1xf32>
    %3 = "tf_device.cluster_func"(%1) <{func = @write_chain_func}> {_replication_info = "cluster__train_helper", _xla_compile_device_type = "TPU", num_cores_per_replica = 1 : i64} : (tensor<1xf32>) -> tensor<1xf32>
    "tf.AssignVariableOp"(%arg0, %3) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1xf32>>>, tensor<1xf32>) -> ()
    tf_device.return
  }
  func.return
}