gfn.utils.prob_calculations
===========================

.. py:module:: gfn.utils.prob_calculations


Functions
---------

.. autoapisummary::

   gfn.utils.prob_calculations.get_trajectory_pbs
   gfn.utils.prob_calculations.get_trajectory_pfs
   gfn.utils.prob_calculations.get_trajectory_pfs_and_pbs
   gfn.utils.prob_calculations.get_transition_pbs
   gfn.utils.prob_calculations.get_transition_pfs
   gfn.utils.prob_calculations.get_transition_pfs_and_pbs


Module Contents
---------------

.. py:function:: get_trajectory_pbs(pb, trajectories, **policy_kwargs)

   Calculate PB log‑probabilities for trajectories.

   Non‑vectorized (per‑step) evaluation with with alignment
       (action at ``t`` with state at ``t+1``) and mask
       ``~is_sink_state[t+1] & ~is_initial_state[t+1] & ~is_dummy[t] & ~is_exit[t]``;
       skip ``t==0``. is supported when specifically needed
       (estimator.is_vectorized=False).

   :param pb: Backward policy estimator, or ``None`` for trees (PB=1).
   :param trajectories: Trajectories to evaluate.
   :param \*\*policy_kwargs: Extra kwargs for ``to_probability_distribution``.

   :returns: ``log_pb`` of shape ``(T, N)``.

   :raises ValueError: If backward trajectories are provided.


.. py:function:: get_trajectory_pfs(pf, trajectories, recalculate_all_logprobs = True, **policy_kwargs)

   Calculate PF log‑probabilities for trajectories.

   Non‑vectorized (per‑step) evaluation with masks
   ``~is_sink_state[t] & ~is_dummy[t]`` & no action‑id indexing is supported when
   specifically needed (estimator.is_vectorized=False).

   :param pf: Forward policy estimator.
   :param trajectories: Trajectories to evaluate.
   :param recalculate_all_logprobs: If True, recompute PF even if cached. Useful for
                                    off-policy training.
   :param \*\*policy_kwargs: Extra kwargs for ``to_probability_distribution``.

   :returns: ``log_pf`` of shape ``(T, N)``.

   :raises ValueError: If backward trajectories are provided.


.. py:function:: get_trajectory_pfs_and_pbs(pf, pb, trajectories, recalculate_all_logprobs = True, **policy_kwargs)

   Calculate PF and PB log‑probabilities for trajectories.

   Delegates to ``get_trajectory_pfs`` and ``get_trajectory_pbs`` while
   forwarding policy kwargs.

   :param pf: Forward policy estimator.
   :param pb: Backward policy estimator, or ``None`` for trees (PB=1).
   :param trajectories: Trajectories to evaluate.
   :param recalculate_all_logprobs: If True, recompute PF even if cached.
   :param \*\*policy_kwargs: Extra kwargs for ``to_probability_distribution``.

   :returns: ``(log_pf[T,N], log_pb[T,N])``


.. py:function:: get_transition_pbs(pb, transitions, **policy_kwargs)

   Calculate PB log‑probabilities for transitions.

   :param pb: Backward policy estimator, or ``None`` for trees (PB=1).
   :param transitions: Transitions to evaluate.
   :param \*\*policy_kwargs: Extra kwargs for ``to_probability_distribution``.

   :returns: ``log_pb`` of shape ``(M,)``.


.. py:function:: get_transition_pfs(pf, transitions, recalculate_all_logprobs = True, **policy_kwargs)

   Calculate PF log‑probabilities for transitions.

   :param pf: Forward policy estimator.
   :param transitions: Transitions to evaluate.
   :param recalculate_all_logprobs: If True, recompute PF even if cached. Useful for
                                    off-policy training.
   :param \*\*policy_kwargs: Extra kwargs for ``to_probability_distribution``.

   :returns: ``log_pf`` of shape ``(M,)``.


.. py:function:: get_transition_pfs_and_pbs(pf, pb, transitions, recalculate_all_logprobs = True, **policy_kwargs)

   Calculate PF and PB log‑probabilities for transitions.

   :param pf: Forward policy estimator.
   :param pb: Backward policy estimator, or ``None`` for trees (PB=1).
   :param transitions: Transitions to evaluate.
   :param recalculate_all_logprobs: If True, recompute PF even if cached. Useful for
                                    off-policy training.
   :param \*\*policy_kwargs: Extra kwargs for ``to_probability_distribution``.

   :returns: ``(log_pf[M], log_pb[M])``.

   :raises ValueError: If backward transitions are provided.