Registry / ai-ml / gram-newton-schulz

gram-newton-schulz

library0.1.4pypypi✓ verified 83d ago

A fast implementation of the Newton-Schulz algorithm for computing matrix square roots and Gram matrix inverses, with support for JIT-compiled kernels via Quack. Current version 0.1.4, requires Python >=3.10. Released under the Dao-AILab organization, with occasional updates.

ai-ml data

pip install gram-newton-schulz

github ↗package ↗

INSTALL

IMPORT

SIG · GRAM-NEWTON-SCHULZ

gram-newton-schulz

ai-mlpythonv0.1.4

harness data pending

Install & Compatibility

Where this runs

No compatibility data collected yet for this library.

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

GramNewtonSchulz

✓ from gram_newton_schulz import GramNewtonSchulz

✗ import GramNewtonSchulz

Class is inside the package, not a top-level module

NewtonSchulzBase

✓ from gram_newton_schulz import NewtonSchulzBase

Compute matrix square root using Gram-Newton-Schulz algorithm.

import torch from gram_newton_schulz import GramNewtonSchulz # Create a symmetric positive-definite matrix A = torch.randn(4, 4, device='cuda' if torch.cuda.is_available() else 'cpu') A = A @ A.T # Initialize the solver with default settings solver = GramNewtonSchulz() # Compute the matrix square root (X such that X @ X ≈ A) X = solver.sqrt(A) print(X)

Debug

Known issues

gotchaInput matrix must be symmetric positive-definite; non-SPD matrices may cause convergence failure or incorrect results.

fix

Ensure your matrix is symmetric and has positive eigenvalues. Consider adding a small regularization term like `A + 1e-6 * torch.eye(A.shape[0])`.

affects: >=0.1.0

gotchaTorch.compile is required for performance; without it, fallback implementation may be very slow or unsupported.

fix

Enable torch.compile by setting `torch.compile(solver.sqrt)` or using the `compile_kwargs` argument.

affects: >=0.1.0

deprecatedThe `StandardNewtonSchulz` class was merged into `GramNewtonSchulz` in v0.1.0. Old imports will break.

fix

Replace `from gram_newton_schulz import StandardNewtonSchulz` with `from gram_newton_schulz import GramNewtonSchulz`.

affects: <0.1.0

Errors

Common errors & fixes

ModuleNotFoundError: No module named 'gram_newton_schulz'

Library not installed, or installed with wrong name.

fix

Run `pip install gram-newton-schulz`. The correct import string is `from gram_newton_schulz import ...`.

RuntimeError: Expected all tensors to be on the same device, but found at least two devices

Solver is initialized on CPU while input tensor is on GPU, or vice versa.

fix

Ensure both the solver (if it holds state) and input tensors are on the same device. Use `solver = GramNewtonSchulz().to(device)`.

Upgrade

Version history

0.1.4latest on PyPI · released May 2, 2026

Audit

Dependencies

torchrequiredCore tensor operations and compilation

quackrequiredJIT kernel backend (>=0.3.7)

Agent activity

9 hits · last 30 days

node

OpenAI (training)

Resources

githubhttps://github.com/Dao-AILab/gram-newton-schulz ↗

packagehttps://pypi.org/project/gram-newton-schulz/ ↗