experimental/no_u_turn_sampler/nuts.py
def _embed_no_none_gradient_check(value_and_gradients_fn):
  """Wraps value and gradients function to assist with None gradients."""
  @functools.wraps(value_and_gradients_fn)
  def func_wrapped(*args, **kwargs):
    """Wrapped function which checks for None gradients."""
    value, grads = value_and_gradients_fn(*args, **kwargs)
    if any(grad is None for grad in grads):
      raise ValueError("Gradient is None for a state.")
    return value, grads
  return func_wrapped
def _embed_no_none_gradient_check(value_and_gradients_fn):
  @functools.wraps(value_and_gradients_fn)
  def func_wrapped(*args, **kwargs):
    """Wrapped function which checks for None gradients."""
    value, grads = value_and_gradients_fn(*args, **kwargs)
    if any(grad is None for grad in grads):
      raise ValueError("Gradient is None for a state.")
    return value, grads
  return func_wrapped
def _embed_no_none_gradient_check ( value_and_gradients_fn ) : @ functools . wraps ( value_and_gradients_fn ) def func_wrapped ( * args , * * kwargs ) : """Wrapped function which checks for None gradients.""" value , grads = value_and_gradients_fn ( * args , * * kwargs ) if any ( grad is None for grad in grads ) : raise ValueError ( "Gradient is None for a state." ) return value , grads return func_wrapped
Wraps value and gradients function to assist with None gradients.
Wraps value and gradients function to assist with None gradients .