nnx.Swish, jax.swish,... change the input shape
Open
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
### System information
- Colab
- Latest version (!pip install git+https://github.com/google/flax.git)
### Problem you have encountered:
print(xBC.shape)
xBC = jax.nn.swish(x) or xBC = nnx.swish(x)
print(xBC.shape)
output shape is not the same as expected
### What you expected to happen:
Output:
(1, 128, 288)
(1, 128, 128)
Expect:
(1, 128, 288)
(1, 128, 288)
Contributor guide
Assessment
This issue has not been assessed yet.