mlx_memory (mlx v0.2.0)

View Source

Summary

Functions

Functions

accumulate_gradients(AccumulatorId, ParameterName, Gradient)

activation_checkpointing(Model, CheckpointLayers)

activation_checkpointing(Model, CheckpointLayers, Options)

batch_evaluations(DeferredComputations)

cache_management(CacheId, Strategy)

chunked_operation(Operation, Tensor, ChunkSize)

chunked_operation(Operation, Tensor, ChunkSize, Options)

copy_on_write(Tensor)

create_gradient_accumulator(AccumulationSteps)

create_gradient_accumulator(AccumulationSteps, Options)

defer_computation(ComputationFun)

force_evaluation(Tensors)

free_memory(_)

get_accumulated_gradients(AccumulatorId)

inplace_operation(Operation, Tensor)

inplace_operation(Operation, Tensor, Options)

memory_compression(Tensor, CompressionLevel)

memory_decompression(CompressedTensor)

memory_efficient_dataloader(Dataset, BatchSize)

memory_efficient_dataloader(Dataset, BatchSize, Options)

memory_leak_detection()

memory_mapping(FilePath, Size)

memory_mapping(FilePath, Size, Options)

memory_optimization_suggestions()

memory_pool_create(PoolSize)

memory_pool_destroy(PoolId)

memory_usage_analysis()

memory_usage_analysis(Options)

optimize_lazy_evaluation(ComputationGraph)

optimize_memory_layout(Tensors)

optimize_memory_layout(Tensors, Options)

preallocate_memory(PoolId, Size)

prefetch_data(DataLoaderState, NumBatches)

recompute_activations(CheckpointedLayers, Inputs)

reset_accumulator(AccumulatorId)

scale_accumulated_gradients(AccumulatorId, ScaleFactor)

shared_memory_tensors(TensorList, SharedMemoryId)

streaming_operation(Operation, InputTensor, OutputSize)

streaming_operation(Operation, InputTensor, OutputSize, Options)