Skip to content

Utilities

jax_morph.ad_utils.safe_norm #

safe_norm(x, axis=-1, keepdims=False)

Compute a Euclidean norm with value and gradient zero at the zero vector.

Parameters:

  • x –

    Input array.

  • axis –

    Axis or axes over which to compute the norm. Defaults to -1.

  • keepdims –

    Whether reduced axes remain with size one. Defaults to False.

Returns:

  • –

    The norm of x with the requested reduced shape.

jax_morph.ad_utils.safe_log #

safe_log(x)

Compute a log that stays finite with a finite gradient for nonpositive inputs.

Parameters:

  • x –

    Input values.

Returns:

  • –

    log(x) where x > 0, and -1e30 otherwise.

jax_morph.ad_utils.safe_divide #

safe_divide(a, b)

Divide two arrays with zero value and gradient where the denominator is zero.

Parameters:

  • a –

    Numerator values.

  • b –

    Denominator values, broadcastable with a.

Returns:

  • –

    a / b where b is nonzero, and zero otherwise.

jax_morph.ad_utils.straight_through #

straight_through(hard, soft)

Return hard exactly while differentiating through soft.

Parameters:

  • hard –

    Values to use in the forward pass.

  • soft –

    Shape-compatible surrogate values for the backward pass.

Returns:

  • –

    An array equal to hard whose gradient is that of soft.

jax_morph.ad_utils.heaviside_st #

heaviside_st(x, temperature=1.0)

Apply an exact Heaviside threshold with a sigmoid surrogate gradient.

Parameters:

  • x –

    Values to threshold at zero.

  • temperature –

    Positive sigmoid temperature for the backward pass. Defaults to 1.0.

Returns:

  • –

    Exact zero-or-one threshold values with sigmoid surrogate gradients.

jax_morph.ad_utils.argmax_st #

argmax_st(logits, temperature=1.0, axis=-1)

Select an exact one-hot argmax with a softmax surrogate gradient.

Parameters:

  • logits –

    Unnormalized category logits.

  • temperature –

    Positive softmax temperature for the backward pass. Defaults to 1.0.

  • axis –

    Category axis. Defaults to -1.

Returns:

  • –

    One-hot argmax values with softmax surrogate gradients.

jax_morph.ad_utils.sample_categorical_st #

sample_categorical_st(
    key, logits, temperature=1.0, axis=-1
)

Draw an exact one-hot categorical sample with a softmax surrogate gradient.

Parameters:

  • key –

    JAX PRNG key.

  • logits –

    Unnormalized category logits.

  • temperature –

    Positive softmax temperature for the backward pass. Defaults to 1.0.

  • axis –

    Category axis. Defaults to -1.

Returns:

  • –

    One-hot categorical samples with softmax surrogate gradients.

jax_morph.ad_utils.sample_bernoulli_st #

sample_bernoulli_st(key, p)

Draw an exact Bernoulli sample with an identity surrogate gradient.

Parameters:

  • key –

    JAX PRNG key.

  • p –

    Success probabilities.

Returns:

  • –

    Zero-or-one samples whose derivative with respect to p is one.

jax_morph.ad_utils.bernoulli_logp #

bernoulli_logp(outcome, p)

Elementwise Bernoulli log-density log p(outcome | p) for outcome in {0, 1}.

Pairs with sample_bernoulli_st: recompute p from the state (through the policy params) and score the recorded 0/1 outcome. Uses safe_log so p at 0 or 1 gives a finite value and gradient rather than NaN.

Parameters:

  • outcome –

    Recorded zero-or-one outcomes.

  • p –

    Success probabilities, broadcastable with outcome.

Returns:

  • –

    Elementwise Bernoulli log-probabilities.

jax_morph.ad_utils.categorical_logp #

categorical_logp(onehot, logits, axis=-1)

Categorical log-density: log softmax(logits) selected by the one-hot outcome.

Pairs with sample_categorical_st: recompute logits from the state (through the policy params) and score the recorded one-hot action. Reduces over axis (the class axis).

Parameters:

  • onehot –

    Recorded one-hot outcomes.

  • logits –

    Unnormalized category logits.

  • axis –

    Category axis to reduce. Defaults to -1.

Returns:

  • –

    Log-probabilities with the category axis removed.

jax_morph.ad_utils.gaussian_logp #

gaussian_logp(x, mean, std)

Elementwise Gaussian log-density log N(x | mean, std**2) (std is the deviation).

Pairs with a reparameterized draw x = mean + std * xi: recompute mean/std from the state (through the policy params) and score the recorded value x - the density a Brownian / Langevin step assigns to its displacement. mean and std broadcast against x; std is assumed strictly positive (a real noise scale), so no boundary guard is applied.

Parameters:

  • x –

    Recorded values.

  • mean –

    Gaussian means, broadcastable with x.

  • std –

    Strictly positive standard deviations, broadcastable with x.

Returns:

  • –

    Elementwise Gaussian log-densities.