Stannum
Fusing Taichi into PyTorch
PRs are always welcomed, please see TODOs and issues.
Why Stannum?
In differentiable rendering including neural rendering, rendering algorithms are transferred to the field of computer vision, but some rendering operations (e.g., ray tracing and direct volume rendering) are not easy to be expressed in tensor operations but in kernels. Differentiable kernels of Taichi enables fast, efficient and differentiable implementation of rendering algorithms while tensor operators provides math expressiveness.
Stannum bridges Taichi and PyTorch to have advantage of both kernel-based and operator-based parallelism.
Documentation and Usage
Please see documentation.
Code sample of Tube:
import taichi as ti
import torch
@ti.kernel
def mul(arr: ti.template(), out: ti.template()):
for i in arr:
out[i] = arr[i] * 2.0
if __name__ == "__main__":
ti.init(ti.cpu)
a = torch.ones(10, requires_grad=True)
tube = Tube() \
.register_input_tensor((10,), torch.float32, "arr") \
.register_output_tensor((10,), torch.float32, "out", True) \
.register_kernel(mul, ["arr", "out"]) \
.finish()
out = tube(a)
loss = out.sum()
loss.backward()
assert torch.allclose(out, torch.ones_like(out) * 2)
assert torch.allclose(a.grad, torch.ones_like(a) * 2)
Installation & Dependencies
Install stannum with pip by
python -m pip install stannum
Make sure you have the following installed:
- PyTorch
- latest Taichi
- For performance concerns, we strongly recommend to use Taichi >= 1.1.3 (see Issue #9 for more information)
Bugs & Issues
Please feel free to file issues. If a runtime error occurs from the dependencies of stannum, you may also want to check the upstream breaking change tracker.
1 maintainer
Metadata
1 maintainer
Release files for stannum 0.9.1
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| stannum-0.9.1.tar.gz | 17.8 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| stannum-0.9.1-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 36.2 kB
Release files / stannum-0.9.1.tar.gz
| Download URL | stannum-0.9.1.tar.gz |
|---|---|
| Size | 17.8 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
2120b99ac8c1897b0ef5b1b56950d6073ee59c69ae0675c38303f66f474ea786
|
|
BLAKE2b-256 checksum How to use checksums |
f22012aeee7e3c6b0acb5ae40c5f4de9f864fd521b13b8111c98af6f090652b6
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/4.0.1 CPython/3.11.1
|
Release files / stannum-0.9.1-py3-none-any.whl
| Download URL | stannum-0.9.1-py3-none-any.whl |
|---|---|
| Size | 18.4 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
bb0ab129dc80a6e748f5033e454bbf017c8225951d2e5c136eeaf35af4d0843d
|
|
BLAKE2b-256 checksum How to use checksums |
d646c6a9fd6513acebf78f55152df19cbcc55eaf84dd6ef14ce0fde7858b630c
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/4.0.1 CPython/3.11.1
|
