-
Notifications
You must be signed in to change notification settings - Fork 9
Make BP closer to jax optimizer #135
Make BP closer to jax optimizer #135
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Made a few comments. Main things to do are:
- Bundle 5 functions into a dataclass for better organization/documentation and easier usage.
- Do some renaming.
LGTM once these are addressed
Also if you want the corresponding issue to be automatically linked I think you need to say resolves #124 |
@StannisZhou all your comments have been adressed |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM after addressing minor comments:
- Change
BPContainer
toBeliefPropagation
(or maybePropagationFunctions
?) - Shorten all
bp_container
to justbp
- Use
dataclass
instead ofNamedTuple
- Update docstrings
@@ -1002,7 +1044,14 @@ def get_beliefs(bp_arrays: BPArrays) -> Any: | |||
) | |||
return beliefs | |||
|
|||
return run_bp, get_bp_state, get_beliefs | |||
bp_container = BPContainer( | |||
init=functools.partial(update, None), |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Does this mean init
can only take keyword arguments? If so mention that in the documentation for init
in BeliefPropagation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
this restriction does not apply, why do you think it would?
b539026
to
35d47d3
Compare
Codecov Report
@@ Coverage Diff @@
## master #135 +/- ##
=========================================
Coverage 100.00% 100.00%
=========================================
Files 13 13
Lines 950 967 +17
=========================================
+ Hits 950 967 +17
Continue to review full report at Codecov.
|
Resolves #124
We make
graph.BP
closer to JAX optimizers https://jax.readthedocs.io/en/latest/jax.example_libraries.optimizers.html