Remove reference cycles when constructing distribution objects
self -> _parameters -> self cycles were creating work for Python's garbage collector in training loops, where Distribution objects may be created repeatedly when executing eagerly. This CL just fixes that narrow memory issue; I'm not convinced dict(locals()) is super efficient, so we may want to follow up on that for performance. Adds a few unit tests tests with run_test_in_graph_and_eager_modes(assert_no_eager_garbage=True). It'd be nice to expand this coverage over time. Includes a small test_util simplification to support this (TFP tests don't like reset_default_graph for some reason). Testing for cycles in the TFP repo will need to wait on the Normal changes from the TF repo syncing. PiperOrigin-RevId: 211520394
Loading
Please sign in to comment