fix: matrix_power exponent gradient — clamp subgradient + Python-scalar support

[bruAristimunha]

Context

Follow-up to #26, which added a gradient w.r.t. the exponent parameter in matrix_power.backward. Two issues remained after that merge.

Bug 1 — spurious huge exponent gradient at clamped eigenvalues

For numerical stability, eigenvalues below get_epsilon(dtype, "eigval_power") (~1e-13 for double) are clamped to that threshold. The X-gradient correctly picks subgradient 0 there:

# core.py:412-413 (existing)
s_deriv = exponent * s_clamped.pow(exponent - 1.0)
s_deriv[s <= threshold] = 0   # subgradient 0 at clamped eigenvalues

But the new exponent-gradient kept flowing through s_modified * log(s_safe) = threshold^e * log(threshold). For s = 1e-30, e = -0.5, this produces grad_exponent ≈ -6.18e+07 instead of the expected ~-0.49.

This makes any training that touches near-singular SPDs through matrix_power blow up.

Bug 2 — Python-scalar exponent crashes backward

Several internal call sites pass a Python float:

# modules/liebn.py:289
matrix_power.apply(X, self.theta)   # self.theta is a Python float

PR #26's backward does .reshape_as(exponent) which crashes with TypeError: reshape_as(): argument 'other' must be Tensor, not float. Reproduced on main.

Fix

1. forward: cast exponent to a 0-d tensor on X's device/dtype via torch.as_tensor so the .reshape_as call always sees a tensor.
2. backward: apply the same subgradient-0 convention to exp_g via torch.where(s > threshold, exp_g, 0.0).
3. Skip the exponent-gradient computation (including the n×n matmul U.mT @ grad_output @ U) when ctx.needs_input_grad[1] is False — saves work on the common path where a Python scalar is passed.
4. Docstring updated: exponent : float or torch.Tensor.

Net diff: +25 / −13 lines.

Tests

tests/test_functional.py::test_matrix_power now covers:

• existing gradcheck with 0-d tensor exponent (unchanged)
• backward with a Python-scalar exponent doesn't crash
• exponent gradient stays bounded (< 1) on a near-singular input where the bug previously produced ~6e7

Full suite: 707 passed, 146 skipped (no regressions).

Authors

• @bruAristimunha

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 22fda2bbbb

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[github-actions]

📚 Documentation Preview

📦 Download Documentation Artifact

Download the documentation-html artifact from the workflow run to view the docs locally.

💡 To enable live previews, add a SURGE_TOKEN secret to this repository. See surge.sh for setup instructions.