# AUTOGENERATED FILE! PLEASE DON'T EDIT
from k1lib.callbacks import Callback, Callbacks
import k1lib, time, torch, math, logging, numpy as np, torch.nn as nn
from functools import partial
import matplotlib.pyplot as plt
__all__ = ["Profiler"]
import k1lib.callbacks.profilers as ps
ComputationProfiler = ps.computation.ComputationProfiler
IOProfiler = ps.io.IOProfiler
MemoryProfiler = ps.memory.MemoryProfiler
TimeProfiler = ps.time.TimeProfiler
[docs]@k1lib.patch(Callback.cls)
class Profiler(Callback):
    """Profiles memory, time, and computational complexity of the network. See over
:mod:`k1lib.callbacks.profilers` for more details on each of these profilers"""
    def __init__(self):
        super().__init__(); self.clear(); self.dependsOn=["Recorder"]
[docs]    def clear(self):
        """Clears every child profilers"""
        self._mpCache=None; self._tpCache=None
        self._cpCache=None; self._ioCache=None 
    def _memory(self): # do this to quickly debug, cause if not, Callback will just raise AttributeError on .memory
        if self._mpCache != None: return self._mpCache
        with self.cbs.context():
            self.cbs.append(mp := MemoryProfiler())
            mp.run(); self._mpCache = mp; return mp
    @property
    def memory(self) -> MemoryProfiler:
        """Gets the memory profiler"""
        return self._memory()
    def _computation(self):
        if self._cpCache != None: return self._cpCache
        with self.cbs.context():
            self.cbs.append(cp := ComputationProfiler(self))
            cp.run(); self._cpCache = cp; return cp
    @property
    def computation(self) -> ComputationProfiler:
        """Gets the computation profiler"""
        return self._computation()
    def _time(self):
        if self._tpCache != None: return self._tpCache
        with self.cbs.context():
            self.cbs.append(tp := TimeProfiler())
            tp.run(); self._tpCache = tp; return tp
    @property
    def time(self) -> TimeProfiler:
        """Gets the time profiler"""
        return self._time()
    def _io(self):
        if self._ioCache != None: return self._ioCache
        with self.cbs.context():
            self.cbs.append(io := IOProfiler())
            io.run(); self._ioCache = io; return io
    @property
    def io(self) -> IOProfiler:
        """Gets the IO profiler"""
        return self._io()
    def __repr__(self):
        return f"""{self._reprHead}, can...
- p.memory: to profile module memory requirements
- p.time: to profile module execution times
- p.computation: to estimate module computation
- p.io: to get input and output shapes of 
{self._reprCan}""" 
@k1lib.patch(Callbacks, docs=Profiler)
def withProfiler(self): return self.append(Profiler())