{"id":31534192,"url":"https://github.com/sashakolpakov/dire-jax","last_synced_at":"2025-10-04T05:18:35.162Z","repository":{"id":280746827,"uuid":"837386271","full_name":"sashakolpakov/dire-jax","owner":"sashakolpakov","description":"DImensionality REduction in JAX","archived":false,"fork":false,"pushed_at":"2025-09-21T14:58:49.000Z","size":244644,"stargazers_count":24,"open_issues_count":2,"forks_count":2,"subscribers_count":3,"default_branch":"main","last_synced_at":"2025-09-26T07:01:46.227Z","etag":null,"topics":["cpu","data-science","data-visualization","dimensionality-reduction","embeddings","gpu","jax","machine-learning","pca","random-projection","tpu","tsne","umap","vector-embeddings"],"latest_commit_sha":null,"homepage":"https://arxiv.org/abs/2503.03156","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/sashakolpakov.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2024-08-02T21:14:17.000Z","updated_at":"2025-09-23T06:30:00.000Z","dependencies_parsed_at":"2025-06-05T04:31:05.295Z","dependency_job_id":"952c3388-4650-4ad6-937f-61d6fe7e728c","html_url":"https://github.com/sashakolpakov/dire-jax","commit_stats":null,"previous_names":["sashakolpakov/dire-jax"],"tags_count":6,"template":false,"template_full_name":null,"purl":"pkg:github/sashakolpakov/dire-jax","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sashakolpakov%2Fdire-jax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sashakolpakov%2Fdire-jax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sashakolpakov%2Fdire-jax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sashakolpakov%2Fdire-jax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/sashakolpakov","download_url":"https://codeload.github.com/sashakolpakov/dire-jax/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sashakolpakov%2Fdire-jax/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":278266896,"owners_count":25958733,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-10-04T02:00:05.491Z","response_time":63,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["cpu","data-science","data-visualization","dimensionality-reduction","embeddings","gpu","jax","machine-learning","pca","random-projection","tpu","tsne","umap","vector-embeddings"],"created_at":"2025-10-04T05:18:29.323Z","updated_at":"2025-10-04T05:18:35.157Z","avatar_url":"https://github.com/sashakolpakov.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\n\u003c!-- Logo + Project title --\u003e\n\u003cp align=\"center\"\u003e\n  \u003cimg src=\"images/logo.png\" alt=\"DiRe-JAX logo\" width=\"280\" style=\"margin-bottom:10px;\"\u003e\n\u003c/p\u003e\n\u003cp align=\"center\"\u003e\n  \u003ca href=\"https://opensource.org/licenses/Apache-2.0\"\u003e\n    \u003cimg alt=\"License\" src=\"https://img.shields.io/badge/License-Apache%202.0-blue.svg\"\u003e\n  \u003c/a\u003e\n  \u003ca href=\"https://www.python.org/downloads/\"\u003e\n    \u003cimg alt=\"Python 3.8+\" src=\"https://img.shields.io/badge/python-3.8+-blue.svg\"\u003e\n  \u003c/a\u003e\n  \u003ca href=\"https://pypi.org/project/dire-jax/\"\u003e\n    \u003cimg alt=\"PyPI\" src=\"https://img.shields.io/pypi/v/dire-jax.svg\"\u003e\n  \u003c/a\u003e\n\u003ca style=\"border-width:0\" href=\"https://doi.org/10.21105/joss.08264\"\u003e\n  \u003cimg src=\"https://joss.theoj.org/papers/10.21105/joss.08264/status.svg\" alt=\"DOI badge\" \u003e\n\u003c/a\u003e\n\u003c/p\u003e\n\u003cp align=\"center\"\u003e\n  \u003ca href=\"https://pepy.tech/projects/dire-jax\"\u003e\n    \u003cimg src=\"https://static.pepy.tech/personalized-badge/dire-jax?period=total\u0026units=ABBREVIATION\u0026left_color=GREY\u0026right_color=BLUE\u0026left_text=downloads\" alt=\"PyPI Downloads\"\u003e\n  \u003c/a\u003e\n  \u003ca href=\"https://github.com/sashakolpakov/dire-jax/actions/workflows/pylint.yml\"\u003e\n    \u003cimg alt=\"CI\" src=\"https://img.shields.io/github/actions/workflow/status/sashakolpakov/dire-jax/pylint.yml?branch=main\u0026label=CI\u0026logo=github\"\u003e\n  \u003c/a\u003e\n  \u003ca href=\"https://github.com/sashakolpakov/dire-jax/actions/workflows/deploy_docs.yml\"\u003e\n    \u003cimg alt=\"Docs\" src=\"https://img.shields.io/github/actions/workflow/status/sashakolpakov/dire-jax/deploy_docs.yml?branch=main\u0026label=Docs\u0026logo=github\"\u003e\n  \u003c/a\u003e\n  \u003ca href=\"https://sashakolpakov.github.io/dire-jax/\"\u003e\n    \u003cimg alt=\"Docs Live\" src=\"https://img.shields.io/website-up-down-green-red/https/sashakolpakov.github.io/dire-jax?label=API%20Documentation\"\u003e\n  \u003c/a\u003e\n\u003c/p\u003e\n\n\n### A high-performance DImensionality REduction package with JAX\n\nDiRe offers fast dimensionality reduction preserving the global dataset structure, with benchmarks showing competitive performance against UMAP and t-SNE. Built with JAX for efficient computation on CPUs and GPUs.\n\n### Quick start\n\n**Basic installation (JAX backend only):**\n```bash    \npip install dire-jax\n```\n\n**With utilities for benchmarking:**\n```bash\npip install dire-jax[utils]\n```\n\n**Complete installation with utilities:**\n```bash\npip install dire-jax[all]\n```\n\n\u003e **Note**: For GPU or TPU acceleration, JAX needs to be specifically installed with hardware support. See the [JAX documentation](https://github.com/google/jax#installation) for more details on enabling GPU/TPU support.\n\n\n**Example usage:**\n```python\nfrom dire_jax import DiRe\nfrom sklearn.datasets import make_blobs\n``` \n\n```python\nn_samples  = 100_000\nn_features = 1_000\nn_centers  = 12\nfeatures_blobs, labels_blobs = make_blobs(n_samples=n_samples, n_features=n_features, centers=n_centers, random_state=42)\n\nreducer_blobs = DiRe(n_components=2,\n                     n_neighbors=16,\n                     init='pca',\n                     max_iter_layout=32,\n                     min_dist=1e-4,\n                     spread=1.0,\n                     cutoff=4.0,\n                     n_sample_dirs=8,\n                     sample_size=16,\n                     neg_ratio=32,\n                     verbose=False,)\n\n_ = reducer_blobs.fit_transform(features_blobs)\nreducer_blobs.visualize(labels=labels_blobs, point_size=4)\n\n```\n\nThe output should look similar to\n\n![12 blobs with 100k points in 1k dimensions embedded in dimension 2](images/blobs_layout.png)\n\n### Documentation \n\nPlease refer to the DiRe API [documentation](https://sashakolpakov.github.io/dire-jax/) for more instructions.\n\n**Project documentation structure:**\n- `/docs/` - API documentation and architecture details\n- `/benchmarking/` - Performance benchmarks and scaling results  \n- `/examples/` - Example usage and demos\n- `/tests/` - Test suite and benchmarking notebooks \n\n### Working paper\n\nOur working paper is available on the arXiv. [![Paper](https://img.shields.io/badge/arXiv-read%20PDF-b31b1b.svg)](https://arxiv.org/abs/2503.03156)\n\n Also, check out the Jupyter notebook with comprehensive benchmarking results and performance analysis. [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](\n  https://colab.research.google.com/github/sashakolpakov/dire-jax/blob/main/benchmarking/dire_benchmarks.ipynb\n)\n\n\n### Performance Characteristics\n\nDiRe-JAX is optimized for small-medium datasets (\u003c50K points) with excellent CPU performance and GPU acceleration via JAX. Features include:\n\n- **Fully vectorized computation** with JIT compilation for optimal performance\n- **Memory-efficient chunking** to handle large datasets without excessive memory usage\n- **Mixed precision arithmetic (MPA)** support for improved performance on modern hardware\n- **Optimized kernel caching** to avoid recompilation and improve runtime efficiency\n- **Large dataset mode** with automatic memory management for datasets \u003e65K points\n\n### Benchmarking and utilities\n\nFor benchmarking utilities and quality metrics:\n```bash\npip install dire-jax[utils]\n```\n\nThis provides access to dimensionality reduction quality metrics and benchmarking routines. Some utilities use external packages for persistent homology computations which may increase runtime. \n\n### Contributing\n\nPlease follow the [contibuting guide](https://sashakolpakov.github.io/dire-jax/contributing.html). Thanks!\n\n### Citation\n\nIf you use this work, please cite it as:\n\n**BibTeX:**\n```bibtex\n@misc{kolpakov-rivin-2025dimensionality,\n  title={Dimensionality reduction for homological stability and global structure preservation},\n  author={Kolpakov, Alexander and Rivin, Igor},\n  year={2025},\n  eprint={2503.03156},\n  archivePrefix={arXiv},\n  primaryClass={cs.LG},\n  url={https://arxiv.org/abs/2503.03156}\n}\n```\n\n**APA Style:**\n```\nKolpakov, A., \u0026 Rivin, I. (2025). Dimensionality reduction for homological stability and global structure preservation. arXiv preprint arXiv:2503.03156. https://arxiv.org/abs/2503.03156\n```\n\n### Acknowledgement\n\nThis work is supported by the Google Cloud Research Award number GCP19980904.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsashakolpakov%2Fdire-jax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fsashakolpakov%2Fdire-jax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsashakolpakov%2Fdire-jax/lists"}