logix-project / logix

AI Logging for Interpretability and Explainability🔬
Apache License 2.0
74 stars 6 forks source link

Update logging setup interface #106

Closed sangkeun00 closed 3 months ago

sangkeun00 commented 3 months ago

This is a proposal to clean up the logging setup interface. Previously, logix.setup used to have three keys of log, save, and statistic, and within statistic we additionally forward, backward, and grad. In fact, we only need three keys of forward, backward, and grad as we show below:

from logix.statistic import Log, Covariance

# before
run.setup({"log": "grad", "save": "grad", "statistic": {"forward": [], "backward": [], "grad": [Covariance]}})

# after
run.setup({"forward": [], "bakcward": [], "grad": [Log, Covariance]})

An additional change we introduce is separating save from the logging setup. Instead users can enable saving both globally (dataset-level) or locally (batch-level) as below:

# global
run.save(True) # run.save(False)

# local
with run(data_id=data_id, save=True):
    ...

I believe this interface would ease users to more easily define and plug-in their own statistics computation primitives.