search

cgarciae / treex

A Pytree Module system for Deep Learning in JAX
https://cgarciae.github.io/treex/
MIT License
214 stars 17 forks source link

Add conv_transpose #64

Closed lkhphuc closed 2 years ago

lkhphuc commented 2 years ago

Add ConvTranspose wrapper. And in Conv, update the padding docstring and add param_dtype for Conv wrapper. I guess this is a new attributes in Flax.

cgarciae commented 2 years ago

Thanks @lkhphuc, reviewed it and it looks great! Merging.