if err != nil {
		return nil, err
	}
	return &r, nil
}

// NewInMemory creates an in-memory allocator.
func NewInMemory(cidr *net.IPNet) (*Range, error) {
	return New(cidr, func(max int, rangeSpec string, offset int) (allocator.Interface, error) {
		return allocator.NewAllocationMapWithOffset(max, rangeSpec, offset), nil
	})
}

      }

      std::shared_ptr<se::DeviceMemoryAllocator> allocator =
          GetAllocator(ctx->device(), GetStream(ctx), platform_info);
      XlaComputationLaunchContext launch_context =
          GetLaunchContext(platform_info, ctx, client, allocator.get());

      const xla::HloInputOutputAliasConfig& input_output_alias =

#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) {

      &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,

    std::optional<std::set<int>> allowed_devices;
  };

  // Creates a new XLA Device.
  XlaDevice(const SessionOptions& session_options, const Options& options);
  ~XlaDevice() override;

  Allocator* GetAllocator(AllocatorAttributes attr) override
      TF_LOCKS_EXCLUDED(mu_);
  void Compute(OpKernel* op_kernel, OpKernelContext* context) override;
  void ComputeAsync(AsyncOpKernel* op_kernel, OpKernelContext* context,

}
def TFL_LSTMKernelTypeAttr : EnumAttr<TFL_Dialect, TFL_LSTMKernelType,
    "lstm_kernel_type_attr">;

def I32ArrayParameter :
    AttrOrTypeParameter<"::llvm::ArrayRef<int32_t>", ""> {
  let allocator = [{$_dst = $_allocator.copyInto($_self);}];
  let cppStorageType = "::llvm::SmallVector<int32_t>";
  let parser = "::mlir::TFL::parseI32Array($_parser)";
  let printer = "$_printer << '[' << $_self << ']'";
}

    TF_CHECK_OK(device_->TryGetDeviceContext(&device_context_));

    // Get the host allocator.
    AllocatorAttributes host_alloc_attr;
    host_alloc_attr.set_on_host(true);
    host_allocator_ = device_->GetAllocator(host_alloc_attr);

    // Get the device allocator. This should give us an AsyncValueAllocator.
    AllocatorAttributes device_alloc_attr;
    device_alloc_attr.set_on_host(false);

  std::shared_ptr<se::DeviceMemoryAllocator> allocator_ptr =
      GetAllocator(ctx->device(), stream, platform_info_);
  se::DeviceMemoryAllocator* allocator = allocator_ptr.get();
  XlaComputationLaunchContext launch_context(
      client, allocator, client->default_device_ordinal(),
      /*allocate_xla_tensors=*/platform_info_.xla_device_metadata() != nullptr,

#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) {

limitations under the License.
*/

package allocator

import (
	"testing"

	"k8s.io/apimachinery/pkg/util/sets"
)

func TestAllocate(t *testing.T) {
	testCases := []struct {
		name      string
		allocator func(max int, rangeSpec string, reserved int) *AllocationBitmap
		max       int
		reserved  int
	}{
		{
			name:      "NewAllocationMap",
			allocator: NewAllocationMapWithOffset,
			max:       32,