Function allocate_mem_buff

codegeex/megatron/memory.py:24–28

Allocate a memory buffer.

(name, numel, dtype, track_usage)

def allocate_mem_buff(name, numel, dtype, track_usage):
    """Allocate a memory buffer."""
    assert name not in _MEM_BUFFS, "memory buffer {} already allocated.".format(name)
    _MEM_BUFFS[name] = MemoryBuffer(name, numel, dtype, track_usage)
    return _MEM_BUFFS[name]


def get_mem_buff(name):