https://github.com/alonfnt/pcax

Differentiable Principal Component Analysis (PCA) implementation in JAX
https://github.com/alonfnt/pcax

differentiable-programming dimensionality-reduction jax pca

Last synced: 3 months ago
JSON representation

Differentiable Principal Component Analysis (PCA) implementation in JAX

• Host: GitHub
• URL: https://github.com/alonfnt/pcax
• Owner: alonfnt
• License: mit
• Created: 2023-05-29T09:15:33.000Z (over 2 years ago)
• Default Branch: main
• Last Pushed: 2025-04-19T07:43:25.000Z (6 months ago)
• Last Synced: 2025-06-25T14:05:30.098Z (3 months ago)
• Topics: differentiable-programming, dimensionality-reduction, jax, pca
• Language: Python
• Homepage:
• Size: 35.2 KB
• Stars: 28
• Watchers: 1
• Forks: 1
• Open Issues: 0
• Metadata Files:
  • Readme: README.md
  • License: LICENSE

Awesome Lists containing this project

README

          


  



[![tests](https://github.com/alonfnt/pcax/actions/workflows/pytest.yml/badge.svg)](https://github.com/alonfnt/pcax/actions/workflows/pytest.yml)

[![PyPI](https://img.shields.io/pypi/v/pcax.svg)](https://pypi.org/project/pcax/)

`pcax` is a minimal PCA implementation in [JAX](https://github.com/jax-ml/jax) that’s both GPU/TPU/CPU‑native and fully differentiable.

It keeps data and computation on-device with zero-copy transfers, lets you backpropagate through your dimensionality reduction step, and plugs directly your JAX workflows for seamless, efficient model integration.

## Usage

```python

import pcax

# Fit the PCA model with 3 components on your data X

state = pcax.fit(X, n_components=3)

# Transform X to its principal components

X_pca = pcax.transform(state, X)

# Recover the original X from its principal components

X_recover = pcax.recover(state, X_pca)

```

## Installation

`pcax` can be installed from PyPI via `pip`

```

pip install pcax

```

## Citation

If you use `pcax` in your research and need to reference it, please cite it as follows:

```

@software{alonso_pcax,

  author = {Alonso, Albert},

  title = {pcax: Minimal Principal Component Analysis (PCA) Implementation in JAX},

  url = {https://github.com/alonfnt/pcax},

  version = {0.1.0}

}

```