Metadata-Version: 2.1
Name: timod
Version: 0.0.17
Summary: A PyTorch module wrapper for Taichi kernels
Home-page: https://github.com/Kiord/timod
Author: Kiord
Author-email: glenn.kerbiriou@gmail.com
License: UNKNOWN
Project-URL: Source, https://github.com/Kiord/timod
Project-URL: Tracker, https://github.com/Kiord/timod/issues
Platform: UNKNOWN
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.7
Description-Content-Type: text/markdown
Requires-Dist: torch>=2.0
Requires-Dist: taichi>=1.6

# timod

TaichIMODule, a tiny pytorch module wrapper for differentiable Taichi kernels.

# Install

`pip install timod`

## Usage

Given a taichi kernel `my_kernel` with input and output tensors typed `ti.types.ndarray()`:

```python
from timod import TaichiKernelModule

# Output tensors specifications (since they will be created on the fly)
output_specs = {
    'out1':(<shape>, <dtype>, <needs gradients ?>), 
    'out2':(<shape>, <dtype>, <needs gradients ?>)
    ...}

# Module creation
taichi_module = TaichiKernelModule(my_kernel, output_specs)

# Kernel input args (torch tensors and scalars)
in1 = ...
in2 = ...
...

# Module call
out1, out2, ... = taichi_module(in1, in2, ...) # kwargs of the taichi kernel inputs are supported 

# Loss function
loss = my_loss(out1, out2, ...)

# Gradient computation
loss.backward()

# Gradient descent
lr = 1e-3
in1 -= lr * in1.grad
in2 -= lr * in2.grad
...

```



