// Permutation from the OIHW (== (output features, input features, height,
// width)) tensor format to HWIO. This is commonly used to transpose convolution
// weights represented as OIHW format to HWIO, which is more desirable for
// certain downstream optimization passes (e.g. XLA).
inline constexpr std::array<int64_t, 4> kOihwToHwioPermutation = {2, 3, 1, 0};

// Returns true if the value has static shape.

// PlatformDeviceId if a mapping between between TfDeviceId and PlatformDeviceId
// is set; otherwise returns the id in the device name. Please refer to
// tensorflow/tsl/framework/device_id.h for more details.
// For mobile or slim build, returns the id in the device name.
TF_CAPI_EXPORT extern int TF_GetDeviceId(TF_OpKernelContext* ctx);

// Returns the Device Name of the device that the context possesses.
//

      // TODO(b/172517537): check if other tensors should go through this
      // check too.
      op.emitError() << "Input tensor [" << input_index
                     << "] is a state tensor, but has more than one use.";
      return failure();
    }
    auto stats = mlir::dyn_cast<DenseFPElementsAttr>(stats_op.getLayerStats());
    if (!stats || stats.getNumElements() != 2) {

#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/local_client.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"

namespace tensorflow {

#include "tensorflow/c/tensor_interface.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/numeric_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/platform.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/tstring.h"

#include <unordered_set>
#include <vector>

#include "tensorflow/core/config/flag_defs.h"
#include "tensorflow/core/config/flags.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
#include "tensorflow/core/lib/gtl/cleanup.h"
#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
#include "tensorflow/core/platform/errors.h"

    if (!expand_op.getResult().hasOneUse()) {
      return rewriter.notifyMatchFailure(
          expand_op, "result for current op has more than 1 use");
    }
    if (!squeeze_op.getResult().hasOneUse()) {
      return rewriter.notifyMatchFailure(
          squeeze_op, "result for current op has more than 1 use");
    }
    // Make sure that the axis in `expand_op` is constant.
    if (auto const_op =

#include "xla/pjrt/pjrt_client.h"
#include "xla/service/hlo.pb.h"
#include "xla/util.h"
#include "tensorflow/core/framework/device.h"
#include "tensorflow/core/lib/strings/proto_serialization.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {

#include <optional>
#include <string>
#include <vector>

#include "absl/container/flat_hash_set.h"
#include "absl/types/optional.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/util/command_line_flags.h"

namespace tensorflow {

struct XlaAutoJitFlag {
  // Control compilation of operators into XLA computations on CPU and GPU

#include "xla/pjrt/pjrt_client.h"
#include "xla/service/shaped_buffer.h"
#include "xla/stream_executor/device_memory_allocator.h"
#include "tensorflow/core/framework/allocation_description.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/thread_annotations.h"

namespace tensorflow {

// Creates a list of updated resource variables.