Chen Ao

`jax.lax.dot` with `preferred_element_type` gives wrong result

3

The bug can be reproduced by the following code ``` import jax.numpy as jnp from jax.random import PRNGKey, normal from jax.lax import dot from jax.config import config config.update("jax_enable_x64", True) A...

[Feature Request] Allow to specify `Conv` padding by strings

9

Hello! I hope to specify the padding in `Conv` by strings like `'SAME'` or `'VALID'`. This is compatible with jax API since it's already implemented in `jax.lax.conv_general_dilated`. Furthermore, I also...