Skip to content

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

Merged
chaoming0625 merged 5 commits intomainfrom
update
Dec 14, 2024
Merged

support unit-aware sparse computation: CSR, CSC, COO#80
chaoming0625 merged 5 commits intomainfrom
update

Conversation

@chaoming0625
Copy link
Member

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
@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