├── .gitignore ├── Makefile ├── license ├── readme.md ├── setup.py ├── tests.py ├── torch_sparse_solve.cpp └── torch_sparse_solve.py /.gitignore: -------------------------------------------------------------------------------- 1 | build/ 2 | dist/ 3 | __pycache__/ 4 | *.egg-info/ 5 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | dist: 2 | python setup.py build sdist bdist_wheel 3 | 4 | build: 5 | python setup.py build_ext --inplace 6 | 7 | clean: 8 | find . -name "dist" | xargs rm -rf 9 | find . -name "build" | xargs rm -rf 10 | find . -name "builds" | xargs rm -rf 11 | find . -name "__pycache__" | xargs rm -rf 12 | find . -name "*.so" | xargs rm -rf 13 | find . -name "*.egg-info" | xargs rm -rf 14 | find . -name ".ipynb_checkpoints" | xargs rm -rf 15 | find . -name ".pytest_cache" | xargs rm -rf 16 | 17 | -------------------------------------------------------------------------------- /license: -------------------------------------------------------------------------------- 1 | GNU LESSER GENERAL PUBLIC LICENSE 2 | Version 2.1, February 1999 3 | 4 | Copyright (C) 1991, 1999 Free Software Foundation, Inc. 5 | 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA 6 | Everyone is permitted to copy and distribute verbatim copies 7 | of this license document, but changing it is not allowed. 8 | 9 | [This is the first released version of the Lesser GPL. It also counts 10 | as the successor of the GNU Library Public License, version 2, hence 11 | the version number 2.1.] 12 | 13 | Preamble 14 | 15 | The licenses for most software are designed to take away your 16 | freedom to share and change it. By contrast, the GNU General Public 17 | Licenses are intended to guarantee your freedom to share and change 18 | free software--to make sure the software is free for all its users. 19 | 20 | This license, the Lesser General Public License, applies to some 21 | specially designated software packages--typically libraries--of the 22 | Free Software Foundation and other authors who decide to use it. You 23 | can use it too, but we suggest you first think carefully about whether 24 | this license or the ordinary General Public License is the better 25 | strategy to use in any particular case, based on the explanations below. 26 | 27 | When we speak of free software, we are referring to freedom of use, 28 | not price. Our General Public Licenses are designed to make sure that 29 | you have the freedom to distribute copies of free software (and charge 30 | for this service if you wish); that you receive source code or can get 31 | it if you want it; that you can change the software and use pieces of 32 | it in new free programs; and that you are informed that you can do 33 | these things. 34 | 35 | To protect your rights, we need to make restrictions that forbid 36 | distributors to deny you these rights or to ask you to surrender these 37 | rights. These restrictions translate to certain responsibilities for 38 | you if you distribute copies of the library or if you modify it. 39 | 40 | For example, if you distribute copies of the library, whether gratis 41 | or for a fee, you must give the recipients all the rights that we gave 42 | you. You must make sure that they, too, receive or can get the source 43 | code. If you link other code with the library, you must provide 44 | complete object files to the recipients, so that they can relink them 45 | with the library after making changes to the library and recompiling 46 | it. And you must show them these terms so they know their rights. 47 | 48 | We protect your rights with a two-step method: (1) we copyright the 49 | library, and (2) we offer you this license, which gives you legal 50 | permission to copy, distribute and/or modify the library. 51 | 52 | To protect each distributor, we want to make it very clear that 53 | there is no warranty for the free library. Also, if the library is 54 | modified by someone else and passed on, the recipients should know 55 | that what they have is not the original version, so that the original 56 | author's reputation will not be affected by problems that might be 57 | introduced by others. 58 | 59 | Finally, software patents pose a constant threat to the existence of 60 | any free program. We wish to make sure that a company cannot 61 | effectively restrict the users of a free program by obtaining a 62 | restrictive license from a patent holder. Therefore, we insist that 63 | any patent license obtained for a version of the library must be 64 | consistent with the full freedom of use specified in this license. 65 | 66 | Most GNU software, including some libraries, is covered by the 67 | ordinary GNU General Public License. This license, the GNU Lesser 68 | General Public License, applies to certain designated libraries, and 69 | is quite different from the ordinary General Public License. We use 70 | this license for certain libraries in order to permit linking those 71 | libraries into non-free programs. 72 | 73 | When a program is linked with a library, whether statically or using 74 | a shared library, the combination of the two is legally speaking a 75 | combined work, a derivative of the original library. The ordinary 76 | General Public License therefore permits such linking only if the 77 | entire combination fits its criteria of freedom. The Lesser General 78 | Public License permits more lax criteria for linking other code with 79 | the library. 80 | 81 | We call this license the "Lesser" General Public License because it 82 | does Less to protect the user's freedom than the ordinary General 83 | Public License. It also provides other free software developers Less 84 | of an advantage over competing non-free programs. These disadvantages 85 | are the reason we use the ordinary General Public License for many 86 | libraries. However, the Lesser license provides advantages in certain 87 | special circumstances. 88 | 89 | For example, on rare occasions, there may be a special need to 90 | encourage the widest possible use of a certain library, so that it becomes 91 | a de-facto standard. To achieve this, non-free programs must be 92 | allowed to use the library. A more frequent case is that a free 93 | library does the same job as widely used non-free libraries. In this 94 | case, there is little to gain by limiting the free library to free 95 | software only, so we use the Lesser General Public License. 96 | 97 | In other cases, permission to use a particular library in non-free 98 | programs enables a greater number of people to use a large body of 99 | free software. For example, permission to use the GNU C Library in 100 | non-free programs enables many more people to use the whole GNU 101 | operating system, as well as its variant, the GNU/Linux operating 102 | system. 103 | 104 | Although the Lesser General Public License is Less protective of the 105 | users' freedom, it does ensure that the user of a program that is 106 | linked with the Library has the freedom and the wherewithal to run 107 | that program using a modified version of the Library. 108 | 109 | The precise terms and conditions for copying, distribution and 110 | modification follow. Pay close attention to the difference between a 111 | "work based on the library" and a "work that uses the library". The 112 | former contains code derived from the library, whereas the latter must 113 | be combined with the library in order to run. 114 | 115 | GNU LESSER GENERAL PUBLIC LICENSE 116 | TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION 117 | 118 | 0. This License Agreement applies to any software library or other 119 | program which contains a notice placed by the copyright holder or 120 | other authorized party saying it may be distributed under the terms of 121 | this Lesser General Public License (also called "this License"). 122 | Each licensee is addressed as "you". 123 | 124 | A "library" means a collection of software functions and/or data 125 | prepared so as to be conveniently linked with application programs 126 | (which use some of those functions and data) to form executables. 127 | 128 | The "Library", below, refers to any such software library or work 129 | which has been distributed under these terms. A "work based on the 130 | Library" means either the Library or any derivative work under 131 | copyright law: that is to say, a work containing the Library or a 132 | portion of it, either verbatim or with modifications and/or translated 133 | straightforwardly into another language. (Hereinafter, translation is 134 | included without limitation in the term "modification".) 135 | 136 | "Source code" for a work means the preferred form of the work for 137 | making modifications to it. For a library, complete source code means 138 | all the source code for all modules it contains, plus any associated 139 | interface definition files, plus the scripts used to control compilation 140 | and installation of the library. 141 | 142 | Activities other than copying, distribution and modification are not 143 | covered by this License; they are outside its scope. The act of 144 | running a program using the Library is not restricted, and output from 145 | such a program is covered only if its contents constitute a work based 146 | on the Library (independent of the use of the Library in a tool for 147 | writing it). Whether that is true depends on what the Library does 148 | and what the program that uses the Library does. 149 | 150 | 1. You may copy and distribute verbatim copies of the Library's 151 | complete source code as you receive it, in any medium, provided that 152 | you conspicuously and appropriately publish on each copy an 153 | appropriate copyright notice and disclaimer of warranty; keep intact 154 | all the notices that refer to this License and to the absence of any 155 | warranty; and distribute a copy of this License along with the 156 | Library. 157 | 158 | You may charge a fee for the physical act of transferring a copy, 159 | and you may at your option offer warranty protection in exchange for a 160 | fee. 161 | 162 | 2. You may modify your copy or copies of the Library or any portion 163 | of it, thus forming a work based on the Library, and copy and 164 | distribute such modifications or work under the terms of Section 1 165 | above, provided that you also meet all of these conditions: 166 | 167 | a) The modified work must itself be a software library. 168 | 169 | b) You must cause the files modified to carry prominent notices 170 | stating that you changed the files and the date of any change. 171 | 172 | c) You must cause the whole of the work to be licensed at no 173 | charge to all third parties under the terms of this License. 174 | 175 | d) If a facility in the modified Library refers to a function or a 176 | table of data to be supplied by an application program that uses 177 | the facility, other than as an argument passed when the facility 178 | is invoked, then you must make a good faith effort to ensure that, 179 | in the event an application does not supply such function or 180 | table, the facility still operates, and performs whatever part of 181 | its purpose remains meaningful. 182 | 183 | (For example, a function in a library to compute square roots has 184 | a purpose that is entirely well-defined independent of the 185 | application. Therefore, Subsection 2d requires that any 186 | application-supplied function or table used by this function must 187 | be optional: if the application does not supply it, the square 188 | root function must still compute square roots.) 189 | 190 | These requirements apply to the modified work as a whole. If 191 | identifiable sections of that work are not derived from the Library, 192 | and can be reasonably considered independent and separate works in 193 | themselves, then this License, and its terms, do not apply to those 194 | sections when you distribute them as separate works. But when you 195 | distribute the same sections as part of a whole which is a work based 196 | on the Library, the distribution of the whole must be on the terms of 197 | this License, whose permissions for other licensees extend to the 198 | entire whole, and thus to each and every part regardless of who wrote 199 | it. 200 | 201 | Thus, it is not the intent of this section to claim rights or contest 202 | your rights to work written entirely by you; rather, the intent is to 203 | exercise the right to control the distribution of derivative or 204 | collective works based on the Library. 205 | 206 | In addition, mere aggregation of another work not based on the Library 207 | with the Library (or with a work based on the Library) on a volume of 208 | a storage or distribution medium does not bring the other work under 209 | the scope of this License. 210 | 211 | 3. You may opt to apply the terms of the ordinary GNU General Public 212 | License instead of this License to a given copy of the Library. To do 213 | this, you must alter all the notices that refer to this License, so 214 | that they refer to the ordinary GNU General Public License, version 2, 215 | instead of to this License. (If a newer version than version 2 of the 216 | ordinary GNU General Public License has appeared, then you can specify 217 | that version instead if you wish.) Do not make any other change in 218 | these notices. 219 | 220 | Once this change is made in a given copy, it is irreversible for 221 | that copy, so the ordinary GNU General Public License applies to all 222 | subsequent copies and derivative works made from that copy. 223 | 224 | This option is useful when you wish to copy part of the code of 225 | the Library into a program that is not a library. 226 | 227 | 4. You may copy and distribute the Library (or a portion or 228 | derivative of it, under Section 2) in object code or executable form 229 | under the terms of Sections 1 and 2 above provided that you accompany 230 | it with the complete corresponding machine-readable source code, which 231 | must be distributed under the terms of Sections 1 and 2 above on a 232 | medium customarily used for software interchange. 233 | 234 | If distribution of object code is made by offering access to copy 235 | from a designated place, then offering equivalent access to copy the 236 | source code from the same place satisfies the requirement to 237 | distribute the source code, even though third parties are not 238 | compelled to copy the source along with the object code. 239 | 240 | 5. A program that contains no derivative of any portion of the 241 | Library, but is designed to work with the Library by being compiled or 242 | linked with it, is called a "work that uses the Library". Such a 243 | work, in isolation, is not a derivative work of the Library, and 244 | therefore falls outside the scope of this License. 245 | 246 | However, linking a "work that uses the Library" with the Library 247 | creates an executable that is a derivative of the Library (because it 248 | contains portions of the Library), rather than a "work that uses the 249 | library". The executable is therefore covered by this License. 250 | Section 6 states terms for distribution of such executables. 251 | 252 | When a "work that uses the Library" uses material from a header file 253 | that is part of the Library, the object code for the work may be a 254 | derivative work of the Library even though the source code is not. 255 | Whether this is true is especially significant if the work can be 256 | linked without the Library, or if the work is itself a library. The 257 | threshold for this to be true is not precisely defined by law. 258 | 259 | If such an object file uses only numerical parameters, data 260 | structure layouts and accessors, and small macros and small inline 261 | functions (ten lines or less in length), then the use of the object 262 | file is unrestricted, regardless of whether it is legally a derivative 263 | work. (Executables containing this object code plus portions of the 264 | Library will still fall under Section 6.) 265 | 266 | Otherwise, if the work is a derivative of the Library, you may 267 | distribute the object code for the work under the terms of Section 6. 268 | Any executables containing that work also fall under Section 6, 269 | whether or not they are linked directly with the Library itself. 270 | 271 | 6. As an exception to the Sections above, you may also combine or 272 | link a "work that uses the Library" with the Library to produce a 273 | work containing portions of the Library, and distribute that work 274 | under terms of your choice, provided that the terms permit 275 | modification of the work for the customer's own use and reverse 276 | engineering for debugging such modifications. 277 | 278 | You must give prominent notice with each copy of the work that the 279 | Library is used in it and that the Library and its use are covered by 280 | this License. You must supply a copy of this License. If the work 281 | during execution displays copyright notices, you must include the 282 | copyright notice for the Library among them, as well as a reference 283 | directing the user to the copy of this License. Also, you must do one 284 | of these things: 285 | 286 | a) Accompany the work with the complete corresponding 287 | machine-readable source code for the Library including whatever 288 | changes were used in the work (which must be distributed under 289 | Sections 1 and 2 above); and, if the work is an executable linked 290 | with the Library, with the complete machine-readable "work that 291 | uses the Library", as object code and/or source code, so that the 292 | user can modify the Library and then relink to produce a modified 293 | executable containing the modified Library. (It is understood 294 | that the user who changes the contents of definitions files in the 295 | Library will not necessarily be able to recompile the application 296 | to use the modified definitions.) 297 | 298 | b) Use a suitable shared library mechanism for linking with the 299 | Library. A suitable mechanism is one that (1) uses at run time a 300 | copy of the library already present on the user's computer system, 301 | rather than copying library functions into the executable, and (2) 302 | will operate properly with a modified version of the library, if 303 | the user installs one, as long as the modified version is 304 | interface-compatible with the version that the work was made with. 305 | 306 | c) Accompany the work with a written offer, valid for at 307 | least three years, to give the same user the materials 308 | specified in Subsection 6a, above, for a charge no more 309 | than the cost of performing this distribution. 310 | 311 | d) If distribution of the work is made by offering access to copy 312 | from a designated place, offer equivalent access to copy the above 313 | specified materials from the same place. 314 | 315 | e) Verify that the user has already received a copy of these 316 | materials or that you have already sent this user a copy. 317 | 318 | For an executable, the required form of the "work that uses the 319 | Library" must include any data and utility programs needed for 320 | reproducing the executable from it. However, as a special exception, 321 | the materials to be distributed need not include anything that is 322 | normally distributed (in either source or binary form) with the major 323 | components (compiler, kernel, and so on) of the operating system on 324 | which the executable runs, unless that component itself accompanies 325 | the executable. 326 | 327 | It may happen that this requirement contradicts the license 328 | restrictions of other proprietary libraries that do not normally 329 | accompany the operating system. Such a contradiction means you cannot 330 | use both them and the Library together in an executable that you 331 | distribute. 332 | 333 | 7. You may place library facilities that are a work based on the 334 | Library side-by-side in a single library together with other library 335 | facilities not covered by this License, and distribute such a combined 336 | library, provided that the separate distribution of the work based on 337 | the Library and of the other library facilities is otherwise 338 | permitted, and provided that you do these two things: 339 | 340 | a) Accompany the combined library with a copy of the same work 341 | based on the Library, uncombined with any other library 342 | facilities. This must be distributed under the terms of the 343 | Sections above. 344 | 345 | b) Give prominent notice with the combined library of the fact 346 | that part of it is a work based on the Library, and explaining 347 | where to find the accompanying uncombined form of the same work. 348 | 349 | 8. You may not copy, modify, sublicense, link with, or distribute 350 | the Library except as expressly provided under this License. Any 351 | attempt otherwise to copy, modify, sublicense, link with, or 352 | distribute the Library is void, and will automatically terminate your 353 | rights under this License. However, parties who have received copies, 354 | or rights, from you under this License will not have their licenses 355 | terminated so long as such parties remain in full compliance. 356 | 357 | 9. You are not required to accept this License, since you have not 358 | signed it. However, nothing else grants you permission to modify or 359 | distribute the Library or its derivative works. These actions are 360 | prohibited by law if you do not accept this License. Therefore, by 361 | modifying or distributing the Library (or any work based on the 362 | Library), you indicate your acceptance of this License to do so, and 363 | all its terms and conditions for copying, distributing or modifying 364 | the Library or works based on it. 365 | 366 | 10. Each time you redistribute the Library (or any work based on the 367 | Library), the recipient automatically receives a license from the 368 | original licensor to copy, distribute, link with or modify the Library 369 | subject to these terms and conditions. You may not impose any further 370 | restrictions on the recipients' exercise of the rights granted herein. 371 | You are not responsible for enforcing compliance by third parties with 372 | this License. 373 | 374 | 11. If, as a consequence of a court judgment or allegation of patent 375 | infringement or for any other reason (not limited to patent issues), 376 | conditions are imposed on you (whether by court order, agreement or 377 | otherwise) that contradict the conditions of this License, they do not 378 | excuse you from the conditions of this License. If you cannot 379 | distribute so as to satisfy simultaneously your obligations under this 380 | License and any other pertinent obligations, then as a consequence you 381 | may not distribute the Library at all. For example, if a patent 382 | license would not permit royalty-free redistribution of the Library by 383 | all those who receive copies directly or indirectly through you, then 384 | the only way you could satisfy both it and this License would be to 385 | refrain entirely from distribution of the Library. 386 | 387 | If any portion of this section is held invalid or unenforceable under any 388 | particular circumstance, the balance of the section is intended to apply, 389 | and the section as a whole is intended to apply in other circumstances. 390 | 391 | It is not the purpose of this section to induce you to infringe any 392 | patents or other property right claims or to contest validity of any 393 | such claims; this section has the sole purpose of protecting the 394 | integrity of the free software distribution system which is 395 | implemented by public license practices. Many people have made 396 | generous contributions to the wide range of software distributed 397 | through that system in reliance on consistent application of that 398 | system; it is up to the author/donor to decide if he or she is willing 399 | to distribute software through any other system and a licensee cannot 400 | impose that choice. 401 | 402 | This section is intended to make thoroughly clear what is believed to 403 | be a consequence of the rest of this License. 404 | 405 | 12. If the distribution and/or use of the Library is restricted in 406 | certain countries either by patents or by copyrighted interfaces, the 407 | original copyright holder who places the Library under this License may add 408 | an explicit geographical distribution limitation excluding those countries, 409 | so that distribution is permitted only in or among countries not thus 410 | excluded. In such case, this License incorporates the limitation as if 411 | written in the body of this License. 412 | 413 | 13. The Free Software Foundation may publish revised and/or new 414 | versions of the Lesser General Public License from time to time. 415 | Such new versions will be similar in spirit to the present version, 416 | but may differ in detail to address new problems or concerns. 417 | 418 | Each version is given a distinguishing version number. If the Library 419 | specifies a version number of this License which applies to it and 420 | "any later version", you have the option of following the terms and 421 | conditions either of that version or of any later version published by 422 | the Free Software Foundation. If the Library does not specify a 423 | license version number, you may choose any version ever published by 424 | the Free Software Foundation. 425 | 426 | 14. If you wish to incorporate parts of the Library into other free 427 | programs whose distribution conditions are incompatible with these, 428 | write to the author to ask for permission. For software which is 429 | copyrighted by the Free Software Foundation, write to the Free 430 | Software Foundation; we sometimes make exceptions for this. Our 431 | decision will be guided by the two goals of preserving the free status 432 | of all derivatives of our free software and of promoting the sharing 433 | and reuse of software generally. 434 | 435 | NO WARRANTY 436 | 437 | 15. BECAUSE THE LIBRARY IS LICENSED FREE OF CHARGE, THERE IS NO 438 | WARRANTY FOR THE LIBRARY, TO THE EXTENT PERMITTED BY APPLICABLE LAW. 439 | EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR 440 | OTHER PARTIES PROVIDE THE LIBRARY "AS IS" WITHOUT WARRANTY OF ANY 441 | KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE 442 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 443 | PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE 444 | LIBRARY IS WITH YOU. SHOULD THE LIBRARY PROVE DEFECTIVE, YOU ASSUME 445 | THE COST OF ALL NECESSARY SERVICING, REPAIR OR CORRECTION. 446 | 447 | 16. IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN 448 | WRITING WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MAY MODIFY 449 | AND/OR REDISTRIBUTE THE LIBRARY AS PERMITTED ABOVE, BE LIABLE TO YOU 450 | FOR DAMAGES, INCLUDING ANY GENERAL, SPECIAL, INCIDENTAL OR 451 | CONSEQUENTIAL DAMAGES ARISING OUT OF THE USE OR INABILITY TO USE THE 452 | LIBRARY (INCLUDING BUT NOT LIMITED TO LOSS OF DATA OR DATA BEING 453 | RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD PARTIES OR A 454 | FAILURE OF THE LIBRARY TO OPERATE WITH ANY OTHER SOFTWARE), EVEN IF 455 | SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH 456 | DAMAGES. 457 | 458 | END OF TERMS AND CONDITIONS 459 | 460 | How to Apply These Terms to Your New Libraries 461 | 462 | If you develop a new library, and you want it to be of the greatest 463 | possible use to the public, we recommend making it free software that 464 | everyone can redistribute and change. You can do so by permitting 465 | redistribution under these terms (or, alternatively, under the terms of the 466 | ordinary General Public License). 467 | 468 | To apply these terms, attach the following notices to the library. It is 469 | safest to attach them to the start of each source file to most effectively 470 | convey the exclusion of warranty; and each file should have at least the 471 | "copyright" line and a pointer to where the full notice is found. 472 | 473 | 474 | Copyright (C) 475 | 476 | This library is free software; you can redistribute it and/or 477 | modify it under the terms of the GNU Lesser General Public 478 | License as published by the Free Software Foundation; either 479 | version 2.1 of the License, or (at your option) any later version. 480 | 481 | This library is distributed in the hope that it will be useful, 482 | but WITHOUT ANY WARRANTY; without even the implied warranty of 483 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU 484 | Lesser General Public License for more details. 485 | 486 | You should have received a copy of the GNU Lesser General Public 487 | License along with this library; if not, write to the Free Software 488 | Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 489 | USA 490 | 491 | Also add information on how to contact you by electronic and paper mail. 492 | 493 | You should also get your employer (if you work as a programmer) or your 494 | school, if any, to sign a "copyright disclaimer" for the library, if 495 | necessary. Here is a sample; alter the names: 496 | 497 | Yoyodyne, Inc., hereby disclaims all copyright interest in the 498 | library `Frob' (a library for tweaking knobs) written by James Random 499 | Hacker. 500 | 501 | , 1 April 1990 502 | Ty Coon, President of Vice 503 | 504 | That's all there is to it! 505 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # Torch Sparse Solve 2 | 3 | An alternative to `torch.solve` for sparse PyTorch CPU tensors using 4 | the efficient 5 | [KLU algorithm](https://ufdcimages.uflib.ufl.edu/UF/E0/01/17/21/00001/palamadai_e.pdf). 6 | 7 | ## CPU tensors only 8 | 9 | This library is a wrapper around the 10 | [SuiteSparse](https://github.com/DrTimothyAldenDavis/SuiteSparse) KLU 11 | algorithms. This means the algorithm is only implemented for 12 | C-arrays and hence is only available for PyTorch CPU 13 | tensors. However, for large, sparse enough tensors, it might still be 14 | worth doing the GPU→CPU conversion. 15 | 16 | ## Usage 17 | 18 | The `torch_sparse_solve` library provides a single function `solve(A, b)`, which solves for `x` in the **batched matrix × batched matrix** 19 | system `Ax=b` for `torch.float64` tensors (notice the different API in comparison to `torch.solve`): 20 | 21 | ```python 22 | import torch 23 | from torch_sparse_solve import solve 24 | torch.manual_seed(42) 25 | mask = torch.tensor([[[1,0,0],[1,1,0],[0,0,1]]], dtype=torch.float64) 26 | A = (mask * torch.randn(4, 3, 3, dtype=torch.float64)).to_sparse() 27 | b = torch.randn(4, 3, 2, dtype=torch.float64) 28 | x = solve(A, b) 29 | 30 | # compare to torch.solve: 31 | A = A.to_dense() 32 | print( (x - torch.solve(b, A)[0] < 1e-9).all() ) 33 | ``` 34 | 35 | `True` 36 | 37 | ## Caveats 38 | 39 | There are two major caveats you should be aware of when using 40 | `torch_sparse_solve.solve(A, b)`: 41 | 42 | - `A` should be 'dense' in the first dimension, i.e. the batch dimension 43 | should contain as many elements as the batch size. 44 | 45 | - `A` should have the same sparsity pattern for every element in the batch. 46 | If this is not the case, you have two options: 47 | 48 | 1. Create a new sparse matrix with the same sparsity pattern for 49 | every element in the batch by adding zeros to the sparse 50 | representation. 51 | 2. **OR** loop over the batch dimension and solve sequentially, i.e. 52 | with shapes `(1, m, m)` and `(1, m, n)` for each element in `A` and `b` 53 | respectively. 54 | 55 | - `solve` is differentiable, but only for the non-zero elements of `A` (which in most 56 | cases is what you want anyway). 57 | 58 | ## Installation 59 | 60 | The library can be installed with `pip`: 61 | 62 | ```bash 63 | pip install torch_sparse_solve 64 | ``` 65 | 66 | Please note that no pre-built wheels exist. This means that `pip` will 67 | attempt to install the library from source. Make sure you have the 68 | necessary dependencies installed for your OS. 69 | 70 | ## Dependencies 71 | 72 | ### Linux 73 | 74 | On Linux, having PyTorch, scipy and suitesparse installed is often 75 | enough to be able install the library (along with the typical 76 | developer tools for your distribution). Run the following inside a 77 | conda environment: 78 | 79 | ```bash 80 | conda install suitesparse scipy 81 | conda install pytorch -c pytorch 82 | pip install torch_sparse_solve 83 | ``` 84 | 85 | ### Windows 86 | 87 | On Windows, the installation process is a bit more involved as 88 | typically the build dependencies are not installed. To install those, 89 | download **Visual Studio Community 2017** from 90 | [here](https://my.visualstudio.com/Downloads?q=visual%20studio%202017&wt.mc_id=o~msft~vscom~older-downloads). 91 | During installation, go to **Workloads** and select the following 92 | workloads: 93 | 94 | - Desktop development with C++ 95 | - Python development 96 | 97 | Then go to **Individual Components** and select the following 98 | additional items: 99 | 100 | - C++/CLI support 101 | - VC++ 2015.3 v14.00 (v140) toolset for desktop 102 | 103 | Then, download and install Microsoft Visual C++ Redistributable from 104 | [here](https://aka.ms/vs/16/release/vc_redist.x64.exe). 105 | 106 | After these installation steps, run the following commands _inside_ a 107 | **x64 Native Tools Command Prompt for VS 2017**, after activating your 108 | conda environment: 109 | 110 | ```bash 111 | set DISTUTILS_USE_SDK=1 112 | conda install suitesparse scipy 113 | conda install pytorch -c pytorch 114 | pip install torch_sparse_solve 115 | ``` 116 | 117 | ## License & Credits 118 | 119 | © Floris Laporte 2020, LGPL-2.1 120 | 121 | This library was partly based on: 122 | 123 | - [SuiteSparse](https://github.com/DrTimothyAldenDavis/SuiteSparse), LGPL-2.1 124 | - [kagami-c/PyKLU](https://github.com/kagami-c/PyKLU), LGPL-2.1 125 | - [scipy.sparse](https://github.com/scipy/scipy/tree/master/scipy/sparse), BSD-3 126 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | import glob 3 | from setuptools import setup, Extension 4 | from torch.utils import cpp_extension 5 | 6 | import torch_sparse_solve 7 | 8 | libroot = os.path.dirname(os.path.dirname(os.__file__)) 9 | if os.name == "nt": # Windows 10 | suitesparse_lib = os.path.join(libroot, "Library", "lib") 11 | suitesparse_include = os.path.join(libroot, "Library", "include", "suitesparse") 12 | else: # Linux / Mac OS 13 | suitesparse_lib = os.path.join(os.path.dirname(libroot), "lib") 14 | suitesparse_include = os.path.join(os.path.dirname(libroot), "include") 15 | 16 | torch_sparse_solve_cpp = Extension( 17 | name="torch_sparse_solve_cpp", 18 | sources=["torch_sparse_solve.cpp"], 19 | include_dirs=[*cpp_extension.include_paths(), suitesparse_include], 20 | library_dirs=[*cpp_extension.library_paths(), suitesparse_lib], 21 | extra_compile_args=[], 22 | libraries=[ 23 | "c10", 24 | "torch", 25 | "torch_cpu", 26 | "torch_python", 27 | "klu", 28 | "btf", 29 | "amd", 30 | "colamd", 31 | "suitesparseconfig", 32 | ], 33 | language="c++", 34 | ) 35 | 36 | try: 37 | with open("readme.md", "r") as f: 38 | long_description = f.read() 39 | except FileNotFoundError: 40 | long_description = torch_sparse_solve.__doc__ 41 | 42 | setup( 43 | name="torch_sparse_solve", 44 | version=torch_sparse_solve.__version__, 45 | author=torch_sparse_solve.__author__, 46 | author_email="floris.laporte@gmail.com", 47 | description=torch_sparse_solve.__doc__, 48 | long_description=long_description, 49 | long_description_content_type="text/markdown", 50 | url="https://github.com/flaport/torch_sparse_solve", 51 | py_modules=["torch_sparse_solve"], 52 | ext_modules=[torch_sparse_solve_cpp], 53 | cmdclass={"build_ext": cpp_extension.BuildExtension}, 54 | classifiers=[ 55 | "Development Status :: 3 - Alpha", 56 | "Programming Language :: Python :: 3", 57 | "Programming Language :: Python :: 2.7", 58 | "License :: OSI Approved :: GNU Lesser General Public License v2 (LGPLv2)", 59 | "Operating System :: OS Independent", 60 | "Topic :: Scientific/Engineering :: Mathematics", 61 | ], 62 | ) 63 | -------------------------------------------------------------------------------- /tests.py: -------------------------------------------------------------------------------- 1 | """ tests for torch_sparse_solve """ 2 | 3 | import numpy 4 | import torch 5 | import ctypes 6 | import pytest 7 | import scipy.sparse 8 | import torch_sparse_solve 9 | import torch_sparse_solve_cpp 10 | 11 | 12 | @pytest.fixture 13 | def gen(): 14 | """ default pytorch random number generator """ 15 | return torch.Generator(device="cpu").manual_seed(42) 16 | 17 | 18 | @pytest.fixture 19 | def A(gen): 20 | """ default batched sparse tensor """ 21 | mask = torch.tensor([[[1, 0, 0], [1, 1, 0], [0, 0, 1]]], dtype=torch.float64) 22 | A = mask * torch.randn(4, 3, 3, dtype=torch.float64, generator=gen) 23 | Asp = A.to_sparse() 24 | Asp.requires_grad_() 25 | return Asp 26 | 27 | 28 | @pytest.fixture 29 | def b(gen): 30 | """ default target tensor """ 31 | return torch.randn(4, 3, 2, dtype=torch.float64, generator=gen) 32 | 33 | 34 | def test_gradcheck(A, b): 35 | """ check if backward grads are correct """ 36 | torch.autograd.gradcheck(torch_sparse_solve.solve, [A, b], check_sparse_nnz=True, check_undefined_grad=False) 37 | 38 | 39 | def test_result(A, b): 40 | """ confirm result """ 41 | x = torch_sparse_solve.solve(A, b) 42 | b2 = torch.bmm(A.to_dense(), x) 43 | numpy.testing.assert_almost_equal( 44 | b2.data.cpu().numpy(), 45 | b.data.cpu().numpy(), 46 | ) 47 | 48 | 49 | def test_comparison_with_torch_solve(A, b): 50 | """ compare with dense torch.solve """ 51 | x_sparse = torch_sparse_solve.solve(A, b) 52 | x_dense = torch.solve(b, A.to_dense())[0] 53 | numpy.testing.assert_almost_equal( 54 | x_sparse.data.cpu().numpy(), x_dense.data.cpu().numpy() 55 | ) 56 | 57 | 58 | def test_coo_to_csc(): 59 | A = torch.tensor( 60 | [ 61 | [0, 0, 0, 0], 62 | [5, 8, 0, 0], 63 | [0, 0, 3, 0], 64 | [0, 6, 0, 0], 65 | ], 66 | dtype=torch.float64, 67 | device="cpu", 68 | ) 69 | Asp = A.to_sparse() 70 | Ai = Asp._indices()[0].int() 71 | Aj = Asp._indices()[1].int() 72 | Ax = Asp._values() 73 | 74 | Ap, Aj, Ax = torch_sparse_solve_cpp._coo_to_csc(A.size(1), Ai, Aj, Ax) 75 | Ap = Ap.data.cpu().numpy() 76 | Aj = Aj.data.cpu().numpy() 77 | Ax = Ax.data.cpu().numpy() 78 | 79 | a = scipy.sparse.csc_matrix(A.detach().cpu().numpy()) 80 | ap = numpy.asarray(a.indptr, dtype=numpy.int64) 81 | aj = numpy.asarray(a.indices, dtype=numpy.int64) 82 | ax = numpy.asarray(a.data, dtype=numpy.float64) 83 | 84 | assert (Ap == ap).all() 85 | assert (Aj == aj).all() 86 | assert (Ax == ax).all() 87 | 88 | 89 | def test_sparse_solver(A, b): 90 | target = torch.solve(b, A.to_dense())[0][0, :, 0] 91 | result = b[0, :, 0].clone() 92 | Ai = A[0]._indices()[0].int() 93 | Aj = A[0]._indices()[1].int() 94 | Ax = A[0]._values() 95 | Ap, Ai, Ax = torch_sparse_solve_cpp._coo_to_csc(A[0].size(1), Ai, Aj, Ax) 96 | torch_sparse_solve_cpp._klu_solve(Ap, Ai, Ax, result) 97 | assert (target - result < 1e-5).all() 98 | 99 | 100 | if __name__ == "__main__": 101 | pytest.main([__file__]) 102 | -------------------------------------------------------------------------------- /torch_sparse_solve.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | at::Tensor solve_forward(at::Tensor A, at::Tensor b); 7 | std::vector solve_backward(at::Tensor grad, at::Tensor A, at::Tensor b, at::Tensor x); 8 | void _klu_solve(at::Tensor Ap, at::Tensor Ai, at::Tensor Ax, at::Tensor b); 9 | std::vector _coo_to_csc(int ncol, at::Tensor Ai, at::Tensor Aj, at::Tensor Ax); 10 | 11 | at::Tensor solve_forward(at::Tensor A, at::Tensor b) { 12 | int p = at::size(b, 0); 13 | int m = at::size(b, 1); 14 | int n = at::size(b, 2); 15 | at::Tensor bflat = at::clone(at::reshape(at::transpose(b, 1, 2), {p, m*n})); 16 | at::Tensor Ax = at::reshape(A._values(), {p, -1}); 17 | at::Tensor Ai = at::reshape(at::_cast_Int(A._indices()[1]), {p, -1}); 18 | at::Tensor Aj = at::reshape(at::_cast_Int(A._indices()[2]), {p, -1}); 19 | for (int i = 0; i < p; i++) { 20 | std::vector Ap_Ai_Ax = _coo_to_csc(m, Ai[i], Aj[i], Ax[i]); 21 | _klu_solve(Ap_Ai_Ax[0], Ap_Ai_Ax[1], Ap_Ai_Ax[2], bflat[i]); // result will be in bflat 22 | } 23 | return at::transpose(bflat.view({p,n,m}), 1, 2); 24 | } 25 | 26 | std::vector solve_backward(at::Tensor grad, at::Tensor A, at::Tensor b, at::Tensor x) { 27 | at::Tensor gradb = solve_forward(A.transpose(1,2).coalesce(), grad); 28 | int batch_size = A.size(0); 29 | std::vector gradA_list; 30 | for (int i=0; i(); 43 | int* ai = Ai.data_ptr(); 44 | double* ax = Ax.data_ptr(); 45 | double* bb = b.data_ptr(); 46 | klu_symbolic* Symbolic; 47 | klu_numeric* Numeric; 48 | klu_common Common; 49 | klu_defaults(&Common); 50 | Symbolic = klu_analyze(ncol, ap, ai, &Common); 51 | Numeric = klu_factor(ap, ai, ax, Symbolic, &Common); 52 | klu_solve(Symbolic, Numeric, ncol, nb/ncol, bb, &Common); 53 | klu_free_symbolic(&Symbolic, &Common); 54 | klu_free_numeric(&Numeric, &Common); 55 | } 56 | 57 | std::vector _coo_to_csc(int ncol, at::Tensor Ai, at::Tensor Aj, at::Tensor Ax) { 58 | int nnz = at::size(Ax, 0); 59 | at::TensorOptions options = at::TensorOptions().dtype(torch::kInt32).device(at::device_of(Ai)); 60 | at::Tensor Bp = at::zeros(ncol+1, options); 61 | at::Tensor Bi = at::zeros_like(Ai); 62 | at::Tensor Bx = at::zeros_like(Ax); 63 | 64 | int* ai = Ai.data_ptr(); 65 | int* aj = Aj.data_ptr(); 66 | double* ax = Ax.data_ptr(); 67 | 68 | int* bp = Bp.data_ptr(); 69 | int* bi = Bi.data_ptr(); 70 | double* bx = Bx.data_ptr(); 71 | 72 | //compute number of non-zero entries per row of A 73 | for (int n = 0; n < nnz; n++) { 74 | bp[aj[n]] += 1; 75 | } 76 | 77 | //cumsum the nnz per row to get Bp 78 | int cumsum = 0; 79 | int temp = 0; 80 | for(int j = 0; j < ncol; j++) { 81 | temp = bp[j]; 82 | bp[j] = cumsum; 83 | cumsum += temp; 84 | } 85 | bp[ncol] = nnz; 86 | 87 | //write Ai, Ax into Bi, Bx 88 | int col = 0; 89 | int dest = 0; 90 | for(int n = 0; n < nnz; n++) { 91 | col = aj[n]; 92 | dest = bp[col]; 93 | bi[dest] = ai[n]; 94 | bx[dest] = ax[n]; 95 | bp[col] += 1; 96 | } 97 | 98 | int last = 0; 99 | for(int i = 0; i <= ncol; i++) { 100 | temp = bp[i]; 101 | bp[i] = last; 102 | last = temp; 103 | } 104 | 105 | return {Bp, Bi, Bx}; 106 | } 107 | 108 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 109 | m.def("solve_forward", &solve_forward, "solve forward"); 110 | m.def("solve_backward", &solve_backward, "solve backward"); 111 | m.def("_klu_solve", &_klu_solve, "sparse solve"); 112 | m.def("_coo_to_csc", &_coo_to_csc, "COO to CSC"); 113 | } 114 | 115 | -------------------------------------------------------------------------------- /torch_sparse_solve.py: -------------------------------------------------------------------------------- 1 | """ A sparse KLU solver for PyTorch """ 2 | 3 | import torch 4 | 5 | __version__ = "0.0.5" 6 | __author__ = "Floris Laporte" 7 | __all__ = ["solve"] 8 | 9 | 10 | def solve(A, b): 11 | """solve a sparse system Ax = b 12 | 13 | Args: 14 | A (torch.sparse.Tensor[b, m, m]): the sparse matrix defining the system. 15 | b (torch.Tensor[b, m, n]): the target matrix b 16 | 17 | Returns: 18 | x (torch.Tensor[b, m, n]): the initially unknown matrix x 19 | 20 | Note: 21 | 'A' should be 'dense' in the first dimension, i.e. the batch dimension 22 | should contain as many elements as the batch size. 23 | 24 | 'A' should have the same sparsity pattern for every element in the batch. 25 | If this is not the case, you have two options: 26 | 1. Create a new sparse matrix with the same sparsity pattern for 27 | every element in the batch by adding zeros to the sparse 28 | representation. 29 | 2. OR loop over the batch dimension and solve sequentially, i.e. 30 | with shapes (1, m, m) and (1, m, n) for each element in 'A' and 'b' 31 | respectively. 32 | """ 33 | return Solve.apply(A, b) 34 | 35 | 36 | class Solve(torch.autograd.Function): 37 | @staticmethod 38 | def forward(ctx, A, b): 39 | if A.ndim != 3 or (A.shape[1] != A.shape[2]) or not A.is_sparse: 40 | raise ValueError( 41 | "'A' should be a batch of square 2D sparse matrices with shape (b, m, m)." 42 | ) 43 | if b.ndim != 3: 44 | raise ValueError("'b' should be a batch of matrices with shape (b, m, n).") 45 | if not A.dtype == torch.float64: 46 | raise ValueError("'A' should be a sparse float64 tensor.") 47 | if not b.dtype == torch.float64: 48 | raise ValueError("'b' should be a float64 tensor.") 49 | from torch_sparse_solve_cpp import solve_forward 50 | 51 | x = solve_forward(A, b) 52 | ctx.save_for_backward(A, b, x) 53 | return x 54 | 55 | @staticmethod 56 | def backward(ctx, grad): 57 | A, b, x = ctx.saved_tensors 58 | from torch_sparse_solve_cpp import solve_backward 59 | 60 | gradA, gradb = solve_backward(grad, A, b, x) 61 | return gradA, gradb 62 | --------------------------------------------------------------------------------