tensorflow_probability/python/distributions/mixture_same_family.py
def _value_and_batch_jacobian(f, x):
  """Enables uniform interface to value and batch jacobian calculation.

  Works in both eager and graph modes.

  Arguments:
    f: The scalar function to evaluate.
    x: The value at which to compute the value and the batch jacobian.

  Returns:
    A tuple (f(x), J(x)), where J(x) is the batch jacobian.
  """
  if tf.executing_eagerly():
    with tf.GradientTape() as tape:
      tape.watch(x)
      value = f(x)
    batch_jacobian = tape.batch_jacobian(value, x)
  else:
    value = f(x)
    batch_jacobian = gradients.batch_jacobian(value, x)
  return value, batch_jacobian
def _value_and_batch_jacobian(f, x):
  """Enables uniform interface to value and batch jacobian calculation.

  Works in both eager and graph modes.

  Arguments:
    f: The scalar function to evaluate.
    x: The value at which to compute the value and the batch jacobian.

  Returns:
    A tuple (f(x), J(x)), where J(x) is the batch jacobian.
  """
  if tf.executing_eagerly():
    with tf.GradientTape() as tape:
      tape.watch(x)
      value = f(x)
    batch_jacobian = tape.batch_jacobian(value, x)
  else:
    value = f(x)
    batch_jacobian = gradients.batch_jacobian(value, x)
  return value, batch_jacobian
Enables uniform interface to value and batch jacobian calculation.

  Works in both eager and graph modes.

  Arguments:
    f: The scalar function to evaluate.
    x: The value at which to compute the value and the batch jacobian.

  Returns:
    A tuple (f(x), J(x)), where J(x) is the batch jacobian.