Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

support unit-aware sparse computation: CSR, CSC, COO #80

Merged
merged 5 commits into from
Dec 14, 2024
Merged

Conversation

chaoming0625
Copy link
Collaborator

This pull request introduces significant updates to the brainunit package, focusing on the addition of sparse matrix support and enhancements to data type promotion functions. Key changes include the integration of the SparseMatrix class, updates to arithmetic operations to handle sparse matrices, and new utility functions for data type promotion.

Sparse Matrix Support:

  • Added SparseMatrix class in brainunit/_sparse_base.py with various arithmetic operations (__add__, __sub__, __mul__, etc.) and placeholder methods for future implementation.
  • Updated arithmetic operations in brainunit/_base.py to handle instances of SparseMatrix. [1] [2] [3] [4] [5] [6] [7]
  • Added brainunit/sparse/__init__.py to expose the COO, CSR, and CSC classes and their associated functions.
  • Created unit tests for the COO class in brainunit/sparse/coo_test.py to validate matrix operations like matvec, matmul, add, sub, mul, div, mod, pos, neg, and abs.

Data Type Promotion Enhancements:

  • Introduced promote_dtypes function in brainunit/math/_fun_keep_unit.py to promote the data types of input arrays to a common type.
  • Added get_promote_dtypes function in brainunit/math/_fun_remove_unit.py to facilitate data type promotion for input arrays.

Documentation:

  • Updated docs/apis/brainunit.sparse.rst to include documentation for the new sparse data structures and operations.

Import and Export Updates:

  • Included sparse module in brainunit/__init__.py and updated the __all__ list to reflect the new module. [1] [2]
  • Imported SparseMatrix in brainunit/_base.py and added split_mantissa_unit function. [1] [2] [3]
  • Imported promote_dtypes from jax in brainunit/math/_fun_keep_unit.py. [1] [2]

These changes significantly enhance the functionality and versatility of the brainunit package, particularly in handling sparse matrices and promoting data types efficiently.

@chaoming0625 chaoming0625 merged commit edeed6c into main Dec 14, 2024
24 checks passed
@chaoming0625 chaoming0625 deleted the update branch December 14, 2024 06:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant