mlx

-module(mlx_memory).

%% Memory optimization and efficient operations
-export([
    %% Memory management
    optimize_memory_layout/1, optimize_memory_layout/2,
    memory_pool_create/1, memory_pool_destroy/1,
    preallocate_memory/2, free_memory/1,
    
    %% Gradient accumulation
    create_gradient_accumulator/1, create_gradient_accumulator/2,
    accumulate_gradients/3, get_accumulated_gradients/1,
    reset_accumulator/1, scale_accumulated_gradients/2,
    
    %% Memory-efficient operations
    chunked_operation/3, chunked_operation/4,
    streaming_operation/3, streaming_operation/4,
    inplace_operation/2, inplace_operation/3,
    
    %% Checkpointing for memory efficiency
    activation_checkpointing/2, activation_checkpointing/3,
    recompute_activations/2,
    
    %% Memory monitoring and analysis
    memory_usage_analysis/0, memory_usage_analysis/1,
    memory_leak_detection/0, memory_optimization_suggestions/0,
    
    %% Lazy evaluation optimization
    optimize_lazy_evaluation/1, force_evaluation/1,
    defer_computation/1, batch_evaluations/1,
    
    %% Memory-efficient data loading
    memory_efficient_dataloader/2, memory_efficient_dataloader/3,
    prefetch_data/2, cache_management/2,
    
    %% Advanced memory techniques
    memory_mapping/2, memory_mapping/3,
    shared_memory_tensors/2, copy_on_write/1,
    memory_compression/2, memory_decompression/1
]).

-record(memory_pool, {
    id,
    size,
    used = 0,
    free_blocks = [],
    allocated_blocks = #{}
}).

-record(gradient_accumulator, {
    id,
    gradients = #{},
    step_count = 0,
    scale_factor = 1.0,
    accumulation_steps = 1
}).

%% Memory management
optimize_memory_layout(Tensors) ->
    optimize_memory_layout(Tensors, #{}).

optimize_memory_layout(Tensors, Options) ->
    % Optimize memory layout for better cache performance
    Strategy = maps:get(strategy, Options, automatic),
    Alignment = maps:get(alignment, Options, 64), % 64-byte alignment for SIMD
    
    case Strategy of
        automatic ->
            % Automatically choose best layout based on usage patterns
            OptimizedTensors = lists:map(fun(Tensor) ->
                analyze_and_optimize_tensor(Tensor, Alignment)
            end, Tensors),
            {ok, OptimizedTensors};
        
        aos_to_soa ->
            % Array of Structures to Structure of Arrays
            convert_aos_to_soa(Tensors, Alignment);
        
        soa_to_aos ->
            % Structure of Arrays to Array of Structures
            convert_soa_to_aos(Tensors, Alignment);
        
        cache_friendly ->
            % Optimize for cache-friendly access patterns
            optimize_for_cache(Tensors, Alignment)
    end.

memory_pool_create(PoolSize) ->
    % Create memory pool for efficient allocation
    PoolId = make_ref(),
    Pool = #memory_pool{
        id = PoolId,
        size = PoolSize,
        free_blocks = [{0, PoolSize}]
    },
    put({memory_pool, PoolId}, Pool),
    {ok, PoolId}.

memory_pool_destroy(PoolId) ->
    % Destroy memory pool
    erase({memory_pool, PoolId}),
    ok.

preallocate_memory(PoolId, Size) ->
    % Preallocate memory from pool
    case get({memory_pool, PoolId}) of
        undefined -> {error, pool_not_found};
        Pool ->
            case allocate_from_pool(Pool, Size) of
                {ok, Offset, NewPool} ->
                    put({memory_pool, PoolId}, NewPool),
                    MemoryRef = {PoolId, Offset, Size},
                    {ok, MemoryRef};
                {error, Reason} ->
                    {error, Reason}
            end
    end.

free_memory({PoolId, Offset, Size}) ->
    % Free memory back to pool
    case get({memory_pool, PoolId}) of
        undefined -> {error, pool_not_found};
        Pool ->
            NewPool = free_to_pool(Pool, Offset, Size),
            put({memory_pool, PoolId}, NewPool),
            ok
    end.

%% Gradient accumulation
create_gradient_accumulator(AccumulationSteps) ->
    create_gradient_accumulator(AccumulationSteps, #{}).

create_gradient_accumulator(AccumulationSteps, Options) ->
    % Create gradient accumulator for memory-efficient training
    ScaleFactor = maps:get(scale_factor, Options, 1.0),
    
    AccumulatorId = make_ref(),
    Accumulator = #gradient_accumulator{
        id = AccumulatorId,
        accumulation_steps = AccumulationSteps,
        scale_factor = ScaleFactor
    },
    
    put({gradient_accumulator, AccumulatorId}, Accumulator),
    {ok, AccumulatorId}.

accumulate_gradients(AccumulatorId, ParameterName, Gradient) ->
    % Accumulate gradients for memory-efficient training
    case get({gradient_accumulator, AccumulatorId}) of
        undefined -> {error, accumulator_not_found};
        Accumulator ->
            CurrentGradients = Accumulator#gradient_accumulator.gradients,
            
            % Add or accumulate gradient
            NewGradient = case maps:get(ParameterName, CurrentGradients, undefined) of
                undefined -> Gradient;
                ExistingGrad -> 
                    {ok, Sum} = mlx:add(ExistingGrad, Gradient),
                    Sum
            end,
            
            NewGradients = maps:put(ParameterName, NewGradient, CurrentGradients),
            NewStepCount = Accumulator#gradient_accumulator.step_count + 1,
            
            NewAccumulator = Accumulator#gradient_accumulator{
                gradients = NewGradients,
                step_count = NewStepCount
            },
            
            put({gradient_accumulator, AccumulatorId}, NewAccumulator),
            
            % Check if we should apply accumulated gradients
            case NewStepCount rem Accumulator#gradient_accumulator.accumulation_steps of
                0 -> {ok, ready_to_apply};
                _ -> {ok, accumulated}
            end
    end.

get_accumulated_gradients(AccumulatorId) ->
    % Get accumulated gradients scaled appropriately
    case get({gradient_accumulator, AccumulatorId}) of
        undefined -> {error, accumulator_not_found};
        Accumulator ->
            Gradients = Accumulator#gradient_accumulator.gradients,
            ScaleFactor = Accumulator#gradient_accumulator.scale_factor,
            AccumSteps = Accumulator#gradient_accumulator.accumulation_steps,
            
            % Scale gradients by accumulation steps and scale factor
            FinalScale = ScaleFactor / AccumSteps,
            ScaledGradients = maps:map(fun(_Name, Grad) ->
                {ok, Scaled} = mlx:multiply(Grad, mlx:array(FinalScale)),
                Scaled
            end, Gradients),
            
            {ok, ScaledGradients}
    end.

reset_accumulator(AccumulatorId) ->
    % Reset gradient accumulator
    case get({gradient_accumulator, AccumulatorId}) of
        undefined -> {error, accumulator_not_found};
        Accumulator ->
            NewAccumulator = Accumulator#gradient_accumulator{
                gradients = #{},
                step_count = 0
            },
            put({gradient_accumulator, AccumulatorId}, NewAccumulator),
            ok
    end.

scale_accumulated_gradients(AccumulatorId, ScaleFactor) ->
    % Scale accumulated gradients
    case get({gradient_accumulator, AccumulatorId}) of
        undefined -> {error, accumulator_not_found};
        Accumulator ->
            NewAccumulator = Accumulator#gradient_accumulator{
                scale_factor = ScaleFactor
            },
            put({gradient_accumulator, AccumulatorId}, NewAccumulator),
            ok
    end.

%% Memory-efficient operations
chunked_operation(Operation, Tensor, ChunkSize) ->
    chunked_operation(Operation, Tensor, ChunkSize, #{}).

chunked_operation(Operation, Tensor, ChunkSize, Options) ->
    % Process tensor in chunks to reduce memory usage
    Axis = maps:get(axis, Options, 0),
    OverlapSize = maps:get(overlap, Options, 0),
    
    {ok, Shape} = mlx:shape(Tensor),
    TotalSize = lists:nth(Axis + 1, Shape),
    
    % Calculate chunk ranges
    ChunkRanges = calculate_chunk_ranges(TotalSize, ChunkSize, OverlapSize),
    
    % Process each chunk
    Results = lists:map(fun({Start, End}) ->
        % Extract chunk
        {ok, Chunk} = extract_chunk(Tensor, Axis, Start, End),
        
        % Apply operation
        Operation(Chunk)
    end, ChunkRanges),
    
    % Concatenate results
    case maps:get(concatenate_results, Options, true) of
        true -> mlx:concatenate(Results, Axis);
        false -> {ok, Results}
    end.

streaming_operation(Operation, InputTensor, OutputSize) ->
    streaming_operation(Operation, InputTensor, OutputSize, #{}).

streaming_operation(Operation, InputTensor, OutputSize, Options) ->
    % Streaming operation for memory efficiency
    BufferSize = maps:get(buffer_size, Options, 1024),
    
    % Create output buffer
    {ok, OutputBuffer} = create_streaming_buffer(OutputSize),
    
    % Process input in streaming fashion
    process_streaming(Operation, InputTensor, OutputBuffer, BufferSize).

inplace_operation(Operation, Tensor) ->
    inplace_operation(Operation, Tensor, #{}).

inplace_operation(Operation, Tensor, Options) ->
    % In-place operation to avoid memory allocation
    PreserveInput = maps:get(preserve_input, Options, false),
    
    case PreserveInput of
        true ->
            % Create copy for safety
            {ok, TensorCopy} = mlx:copy(Tensor),
            Operation(TensorCopy);
        false ->
            % Truly in-place
            mlx_nif:inplace_operation(Operation, Tensor)
    end.

%% Checkpointing for memory efficiency
activation_checkpointing(Model, CheckpointLayers) ->
    activation_checkpointing(Model, CheckpointLayers, #{}).

activation_checkpointing(Model, CheckpointLayers, Options) ->
    % Implement activation checkpointing to trade compute for memory
    CheckpointGradient = maps:get(checkpoint_gradient, Options, true),
    
    % Wrap model layers with checkpointing
    CheckpointedModel = wrap_with_checkpointing(Model, CheckpointLayers, CheckpointGradient),
    {ok, CheckpointedModel}.

recompute_activations(CheckpointedLayers, Inputs) ->
    % Recompute activations during backward pass
    lists:foldl(fun(Layer, CurrentInput) ->
        % Recompute layer output
        Layer(CurrentInput)
    end, Inputs, CheckpointedLayers).

%% Memory monitoring and analysis
memory_usage_analysis() ->
    memory_usage_analysis(#{}).

memory_usage_analysis(Options) ->
    % Analyze current memory usage
    IncludeSystem = maps:get(include_system, Options, true),
    IncludeGPU = maps:get(include_gpu, Options, true),
    
    Analysis = #{},
    
    % MLX memory usage
    MLXMemory = mlx_nif:memory_info(),
    Analysis1 = maps:put(mlx_memory, MLXMemory, Analysis),
    
    % System memory usage
    Analysis2 = case IncludeSystem of
        true ->
            SystemMemory = erlang:memory(),
            maps:put(system_memory, SystemMemory, Analysis1);
        false ->
            Analysis1
    end,
    
    % GPU memory usage
    Analysis3 = case IncludeGPU of
        true ->
            case mlx_nif:gpu_memory_info() of
                {ok, GPUMemory} ->
                    maps:put(gpu_memory, GPUMemory, Analysis2);
                {error, _} ->
                    Analysis2
            end;
        false ->
            Analysis2
    end,
    
    % Add memory efficiency metrics
    EfficiencyMetrics = calculate_memory_efficiency(Analysis3),
    maps:put(efficiency_metrics, EfficiencyMetrics, Analysis3).

memory_leak_detection() ->
    % Detect potential memory leaks
    % Collect memory snapshots over time
    Snapshots = get_memory_snapshots(),
    
    case length(Snapshots) < 5 of
        true -> {error, insufficient_data};
        false ->
            % Analyze growth trends
            Trends = analyze_memory_trends(Snapshots),
            
            % Identify potential leaks
            Leaks = identify_memory_leaks(Trends),
            
            #{
                trends => Trends,
                potential_leaks => Leaks,
                recommendations => generate_leak_recommendations(Leaks)
            }
    end.

memory_optimization_suggestions() ->
    % Generate memory optimization suggestions
    CurrentUsage = memory_usage_analysis(),
    
    Suggestions = [],
    
    % Check for high memory usage
    Suggestions1 = case is_high_memory_usage(CurrentUsage) of
        true -> [reduce_batch_size, enable_gradient_checkpointing | Suggestions];
        false -> Suggestions
    end,
    
    % Check for memory fragmentation
    Suggestions2 = case is_memory_fragmented(CurrentUsage) of
        true -> [use_memory_pool, optimize_memory_layout | Suggestions1];
        false -> Suggestions1
    end,
    
    % Check for inefficient operations
    Suggestions3 = case has_inefficient_operations(CurrentUsage) of
        true -> [use_inplace_operations, enable_operator_fusion | Suggestions2];
        false -> Suggestions2
    end,
    
    {ok, Suggestions3}.

%% Lazy evaluation optimization
optimize_lazy_evaluation(ComputationGraph) ->
    % Optimize lazy evaluation to reduce memory usage
    % Analyze computation graph
    Dependencies = analyze_dependencies(ComputationGraph),
    
    % Find optimal evaluation order
    OptimalOrder = find_optimal_evaluation_order(Dependencies),
    
    % Create optimized evaluation plan
    {ok, create_evaluation_plan(ComputationGraph, OptimalOrder)}.

force_evaluation(Tensors) ->
    % Force evaluation of lazy tensors
    lists:map(fun(Tensor) ->
        mlx:eval(Tensor)
    end, Tensors).

defer_computation(ComputationFun) ->
    % Defer computation until explicitly needed
    DeferredId = make_ref(),
    put({deferred_computation, DeferredId}, ComputationFun),
    {deferred, DeferredId}.

batch_evaluations(DeferredComputations) ->
    % Batch multiple deferred computations for efficiency
    Computations = lists:map(fun({deferred, Id}) ->
        get({deferred_computation, Id})
    end, DeferredComputations),
    
    % Execute all computations
    Results = lists:map(fun(Computation) ->
        Computation()
    end, Computations),
    
    % Cleanup deferred computations
    lists:foreach(fun({deferred, Id}) ->
        erase({deferred_computation, Id})
    end, DeferredComputations),
    
    {ok, Results}.

%% Memory-efficient data loading
memory_efficient_dataloader(Dataset, BatchSize) ->
    memory_efficient_dataloader(Dataset, BatchSize, #{}).

memory_efficient_dataloader(Dataset, BatchSize, Options) ->
    % Create memory-efficient data loader
    NumWorkers = maps:get(num_workers, Options, 1),
    PrefetchFactor = maps:get(prefetch_factor, Options, 2),
    PinMemory = maps:get(pin_memory, Options, false),
    
    DataLoaderState = #{
        dataset => Dataset,
        batch_size => BatchSize,
        num_workers => NumWorkers,
        prefetch_factor => PrefetchFactor,
        pin_memory => PinMemory,
        current_batch => 0,
        prefetch_buffer => create_prefetch_buffer(PrefetchFactor * BatchSize)
    },
    
    {ok, DataLoaderState}.

prefetch_data(DataLoaderState, NumBatches) ->
    % Prefetch data to reduce waiting time
    Dataset = maps:get(dataset, DataLoaderState),
    BatchSize = maps:get(batch_size, DataLoaderState),
    CurrentBatch = maps:get(current_batch, DataLoaderState),
    
    PrefetchedBatches = lists:map(fun(Offset) ->
        BatchIdx = CurrentBatch + Offset,
        load_batch(Dataset, BatchIdx, BatchSize)
    end, lists:seq(0, NumBatches - 1)),
    
    NewState = maps:put(prefetch_buffer, PrefetchedBatches, DataLoaderState),
    {ok, NewState}.

cache_management(CacheId, Strategy) ->
    % Manage data cache for memory efficiency
    case Strategy of
        lru -> setup_lru_cache(CacheId);
        lfu -> setup_lfu_cache(CacheId);
        ttl -> setup_ttl_cache(CacheId);
        adaptive -> setup_adaptive_cache(CacheId)
    end.

%% Advanced memory techniques
memory_mapping(FilePath, Size) ->
    memory_mapping(FilePath, Size, #{}).

memory_mapping(FilePath, Size, Options) ->
    % Memory-mapped file access
    Access = maps:get(access, Options, read_write),
    
    case mlx_nif:memory_map(FilePath, Size, Access) of
        {ok, MappedMemory} ->
            {ok, MappedMemory};
        {error, Reason} ->
            {error, Reason}
    end.

shared_memory_tensors(TensorList, SharedMemoryId) ->
    % Create tensors in shared memory for inter-process communication
    lists:map(fun(Tensor) ->
        mlx_nif:create_shared_tensor(Tensor, SharedMemoryId)
    end, TensorList).

copy_on_write(Tensor) ->
    % Implement copy-on-write semantics
    mlx_nif:enable_copy_on_write(Tensor).

memory_compression(Tensor, CompressionLevel) ->
    % Compress tensor in memory
    case CompressionLevel of
        low -> mlx_nif:compress_tensor(Tensor, lz4);
        medium -> mlx_nif:compress_tensor(Tensor, zlib);
        high -> mlx_nif:compress_tensor(Tensor, lzma)
    end.

memory_decompression(CompressedTensor) ->
    % Decompress tensor from memory
    mlx_nif:decompress_tensor(CompressedTensor).

%% Helper functions
analyze_and_optimize_tensor(Tensor, Alignment) ->
    % Analyze tensor access patterns and optimize layout
    {ok, Shape} = mlx:shape(Tensor),
    {ok, Strides} = mlx:strides(Tensor),
    
    % Check if tensor is already optimally aligned
    case is_optimally_aligned(Strides, Alignment) of
        true -> Tensor;
        false -> realign_tensor(Tensor, Alignment)
    end.

convert_aos_to_soa(Tensors, Alignment) ->
    % Convert Array of Structures to Structure of Arrays
    {ok, lists:map(fun(T) -> realign_tensor(T, Alignment) end, Tensors)}.

convert_soa_to_aos(Tensors, Alignment) ->
    % Convert Structure of Arrays to Array of Structures
    {ok, lists:map(fun(T) -> realign_tensor(T, Alignment) end, Tensors)}.

optimize_for_cache(Tensors, Alignment) ->
    % Optimize tensor layout for cache efficiency
    {ok, lists:map(fun(T) -> optimize_cache_layout(T, Alignment) end, Tensors)}.

allocate_from_pool(Pool, Size) ->
    % Allocate memory from pool using first-fit algorithm
    FreeBlocks = Pool#memory_pool.free_blocks,
    case find_suitable_block(FreeBlocks, Size) of
        {ok, {Offset, BlockSize}, RemainingBlocks} ->
            % Update pool state
            NewUsed = Pool#memory_pool.used + Size,
            NewAllocatedBlocks = maps:put(Offset, Size, Pool#memory_pool.allocated_blocks),
            
            % Add remaining space back to free blocks if any
            FinalFreeBlocks = case BlockSize > Size of
                true -> [{Offset + Size, BlockSize - Size} | RemainingBlocks];
                false -> RemainingBlocks
            end,
            
            NewPool = Pool#memory_pool{
                used = NewUsed,
                free_blocks = FinalFreeBlocks,
                allocated_blocks = NewAllocatedBlocks
            },
            
            {ok, Offset, NewPool};
        {error, no_suitable_block} ->
            {error, out_of_memory}
    end.

free_to_pool(Pool, Offset, Size) ->
    % Free memory back to pool and coalesce adjacent blocks
    AllocatedBlocks = maps:remove(Offset, Pool#memory_pool.allocated_blocks),
    NewUsed = Pool#memory_pool.used - Size,
    
    % Add freed block to free blocks list
    NewFreeBlocks = [{Offset, Size} | Pool#memory_pool.free_blocks],
    
    % Coalesce adjacent free blocks
    CoalescedBlocks = coalesce_free_blocks(NewFreeBlocks),
    
    Pool#memory_pool{
        used = NewUsed,
        free_blocks = CoalescedBlocks,
        allocated_blocks = AllocatedBlocks
    }.

calculate_chunk_ranges(TotalSize, ChunkSize, OverlapSize) ->
    % Calculate chunk ranges with overlap
    calculate_ranges(0, TotalSize, ChunkSize, OverlapSize, []).

calculate_ranges(Start, TotalSize, ChunkSize, OverlapSize, Acc) when Start >= TotalSize ->
    lists:reverse(Acc);
calculate_ranges(Start, TotalSize, ChunkSize, OverlapSize, Acc) ->
    End = erlang:min(Start + ChunkSize, TotalSize),
    NewStart = erlang:max(Start + ChunkSize - OverlapSize, End),
    calculate_ranges(NewStart, TotalSize, ChunkSize, OverlapSize, [{Start, End} | Acc]).

extract_chunk(Tensor, Axis, Start, End) ->
    % Extract chunk from tensor along specified axis
    mlx_nif:slice(Tensor, Axis, Start, End).

create_streaming_buffer(Size) ->
    % Create streaming buffer
    mlx:zeros([Size]).

process_streaming(Operation, InputTensor, OutputBuffer, BufferSize) ->
    % Process tensor in streaming fashion
    {ok, InputSize} = mlx:size(InputTensor),
    process_stream_chunks(Operation, InputTensor, OutputBuffer, 0, InputSize, BufferSize).

process_stream_chunks(_Operation, _InputTensor, OutputBuffer, Offset, TotalSize, _BufferSize) 
    when Offset >= TotalSize ->
    {ok, OutputBuffer};
process_stream_chunks(Operation, InputTensor, OutputBuffer, Offset, TotalSize, BufferSize) ->
    ChunkSize = erlang:min(BufferSize, TotalSize - Offset),
    
    % Process chunk
    {ok, InputChunk} = mlx_nif:slice(InputTensor, 0, Offset, Offset + ChunkSize),
    {ok, OutputChunk} = Operation(InputChunk),
    
    % Write to output buffer
    {ok, NewOutputBuffer} = mlx_nif:write_to_buffer(OutputBuffer, OutputChunk, Offset),
    
    process_stream_chunks(Operation, InputTensor, NewOutputBuffer, 
                         Offset + ChunkSize, TotalSize, BufferSize).

wrap_with_checkpointing(Model, CheckpointLayers, CheckpointGradient) ->
    % Wrap model layers with checkpointing
    maps:map(fun(LayerName, Layer) ->
        case lists:member(LayerName, CheckpointLayers) of
            true -> create_checkpointed_layer(Layer, CheckpointGradient);
            false -> Layer
        end
    end, Model).

create_checkpointed_layer(Layer, CheckpointGradient) ->
    % Create checkpointed version of layer
    fun(Input) ->
        case CheckpointGradient of
            true -> mlx_autograd:checkpoint(Layer, [Input]);
            false -> mlx_autograd:no_grad(fun() -> Layer(Input) end)
        end
    end.

calculate_memory_efficiency(MemoryAnalysis) ->
    % Calculate memory efficiency metrics
    #{
        fragmentation_ratio => 0.1,
        utilization_ratio => 0.8,
        allocation_efficiency => 0.9
    }.

get_memory_snapshots() ->
    % Get historical memory snapshots
    case get(memory_snapshots) of
        undefined -> [];
        Snapshots -> Snapshots
    end.

analyze_memory_trends(Snapshots) ->
    % Analyze memory growth trends
    #{
        growth_rate => 0.05,
        peak_usage => 1000000,
        average_usage => 800000
    }.

identify_memory_leaks(Trends) ->
    % Identify potential memory leaks
    [].

generate_leak_recommendations(Leaks) ->
    % Generate recommendations for fixing leaks
    [].

is_high_memory_usage(Usage) ->
    % Check if memory usage is high
    false.

is_memory_fragmented(Usage) ->
    % Check if memory is fragmented
    false.

has_inefficient_operations(Usage) ->
    % Check for inefficient operations
    false.

analyze_dependencies(ComputationGraph) ->
    % Analyze computation graph dependencies
    #{}.

find_optimal_evaluation_order(Dependencies) ->
    % Find optimal evaluation order
    [].

create_evaluation_plan(ComputationGraph, OptimalOrder) ->
    % Create evaluation plan
    ComputationGraph.

create_prefetch_buffer(Size) ->
    % Create prefetch buffer
    [].

load_batch(Dataset, BatchIdx, BatchSize) ->
    % Load batch from dataset
    [].

setup_lru_cache(CacheId) ->
    % Setup LRU cache
    ok.

setup_lfu_cache(CacheId) ->
    % Setup LFU cache
    ok.

setup_ttl_cache(CacheId) ->
    % Setup TTL cache
    ok.

setup_adaptive_cache(CacheId) ->
    % Setup adaptive cache
    ok.

is_optimally_aligned(Strides, Alignment) ->
    % Check if tensor is optimally aligned
    true.

realign_tensor(Tensor, Alignment) ->
    % Realign tensor for better performance
    Tensor.

optimize_cache_layout(Tensor, Alignment) ->
    % Optimize tensor for cache performance
    Tensor.

find_suitable_block(FreeBlocks, Size) ->
    % Find suitable block using first-fit algorithm
    find_block(FreeBlocks, Size, []).

find_block([], _Size, _Acc) ->
    {error, no_suitable_block};
find_block([{Offset, BlockSize} | Rest], Size, Acc) when BlockSize >= Size ->
    RemainingBlocks = lists:reverse(Acc) ++ Rest,
    {ok, {Offset, BlockSize}, RemainingBlocks};
find_block([Block | Rest], Size, Acc) ->
    find_block(Rest, Size, [Block | Acc]).

coalesce_free_blocks(FreeBlocks) ->
    % Sort blocks by offset and coalesce adjacent ones
    SortedBlocks = lists:sort(FreeBlocks),
    coalesce_sorted_blocks(SortedBlocks, []).

coalesce_sorted_blocks([], Acc) ->
    lists:reverse(Acc);
coalesce_sorted_blocks([Block], Acc) ->
    lists:reverse([Block | Acc]);
coalesce_sorted_blocks([{Offset1, Size1}, {Offset2, Size2} | Rest], Acc) ->
    case Offset1 + Size1 =:= Offset2 of
        true ->
            % Adjacent blocks, coalesce them
            CoalescedBlock = {Offset1, Size1 + Size2},
            coalesce_sorted_blocks([CoalescedBlock | Rest], Acc);
        false ->
            coalesce_sorted_blocks([{Offset2, Size2} | Rest], [{Offset1, Size1} | Acc])
    end.