fmt.Fprintln(w, "\"internal/unsafeheader\"")
	fmt.Fprintln(w, "\"math/bits\"")
	fmt.Fprintln(w, "\"sync\"")
	fmt.Fprintln(w, "\"unsafe\"")
	fmt.Fprintln(w, ")")
	for _, a := range allocators {
		genAllocator(w, a)
	}
	for _, d := range deriveds {
		for _, base := range allocators {
			if base.name == d.base {
				genDerived(w, d, base)
				break
			}
		}
	}
	// gofmt result
	b := w.Bytes()

    tensorflow::AllocatorAttributes host_alloc_attr;
    host_alloc_attr.set_on_host(true);
    host_allocator_ = device_->GetAllocator(host_alloc_attr);

    tensorflow::AllocatorAttributes device_alloc_attr;
    device_alloc_attr.set_on_host(false);
    device_allocator_ = device_->GetAllocator(device_alloc_attr);

    tensorflow::DeviceContext* device_context;
    auto status = device_->TryGetDeviceContext(&device_context);

    AllocatorAttributes host_alloc_attr;
    host_alloc_attr.set_on_host(true);
    host_allocator_ = device_->GetAllocator(host_alloc_attr);

    AllocatorAttributes device_alloc_attr;
    device_alloc_attr.set_on_host(false);
    device_allocator_ = device_->GetAllocator(device_alloc_attr);
  }

 protected:
  std::unique_ptr<Device> device_;
  Allocator* host_allocator_;
  Allocator* device_allocator_;

      &fbody->record->fdef().signature(), node_def, fbody->arg_types,
      fbody->ret_types);
  OpKernelConstruction construction(DeviceType(dev->device_type()), dev,
                                    dev->GetAllocator(AllocatorAttributes()),
                                    flr, dev->resource_manager(), props,
                                    input_memory_types, output_memory_types,

  options.graph_def_version = function_library.graph_def_version();
  options.allow_cpu_custom_calls =
      (platform_info.platform_id() == se::host::kHostPlatformId);
  options.device_allocator = GetAllocator(device, stream, platform_info);
  if (platform_info.xla_device_metadata()) {
    options.shape_determination_fns =
        platform_info.xla_device_metadata()->default_shape_determination_fns();
  }

#include "tensorflow/core/platform/test.h"

namespace tensorflow {
namespace {

class DummyDevice : public DeviceBase {
 public:
  explicit DummyDevice(Env* env) : DeviceBase(env) {}
  Allocator* GetAllocator(AllocatorAttributes /*attr*/) override {
    return cpu_allocator();
  }
};

void TestBitcastOp(Tensor* input_tensor, DataType out_type,
                   TensorShape expected_shape, error::Code expected_code) {

//
// This is necessary because for XLA devices the underlying TF allocator returns
// dummy tensors.
//
// `stream` parameter is nullable when running on host.
std::shared_ptr<se::DeviceMemoryAllocator> GetAllocator(
    DeviceBase* device, se::Stream* stream,
    const XlaPlatformInfo& platform_info);

}  // namespace tensorflow

#include "tensorflow/core/protobuf/error_codes.pb.h"

namespace tensorflow {
namespace {

class DummyDevice : public DeviceBase {
 public:
  explicit DummyDevice(Env* env) : DeviceBase(env) {}
  Allocator* GetAllocator(AllocatorAttributes /*attr*/) override {
    return cpu_allocator();
  }
};

// Helper for comparing output and expected output
void ExpectSummaryMatches(const Summary& actual, const string& expected_str) {