flax nn.tabulate Incorrectly Reports FLOPs and VJP FLOPs
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): `Ubuntu 22.04.4 LTS x86_64`
- Flax, jax, jaxlib versions:
```
Name: flax
Version: 0.8.4
---
Name: jax
Version: 0.4.30
---
Name: jaxlib
Version: 0.4.30
```
- Python version: `Python 3.12.4`
- GPU/TPU model and memory: ` NVIDIA GeForce GTX 3080 Ti`
- CUDA version: `12.2`
### Problem you have encountered:
When running a script to tabulate the model summary including FLOPs and VJP FLOPs using Flax's [`nn.tabulate`](command:_github.copilot.openSymbolFromReferences?%5B%7B%22%24mid%22%3A1%2C%22path%22%3A%22%2Fhome%2Flab%2Fminiforge3%2Fenvs%2Fjax_base%2Flib%2Fpython3.12%2Fsite-packages%2Fflax%2Flinen%2F__init__.py%22%2C%22scheme%22%3A%22file%22%7D%2C%7B%22line%22%3A0%2C%22character%22%3A0%7D%5D "../../miniforge3/envs/jax_base/lib/python3.12/site-packages/flax/linen/__init__.py") function, the output incorrectly shows both FLOPs and VJP FLOPs as 0. This is unexpected as the model does perform computations that should result in a non-zero FLOPs count, and especially the VJP FLOPs should be a non-zero integer value given the model's structure and operations.
### What you expected to happen:
The expected output should correctly calculate and display the FLOPs and VJP FLOPs for each layer in the model.
### Logs, error messages, etc:
```
import flax.linen as nn
import jax
import jax.numpy as jnp
class Foo(nn.Module):
@nn.compact
def __call__(self, x):
h = nn.Dense(4)(x)
return nn.Dense(2)(h)
x = jnp.ones((16, 9))
tabulate_fn = nn.tabulate(
Foo(), jax.random.PRNGKey(0), compute_flops=True, compute_vjp_flops=True)
print(tabulate_fn(x))
Foo Summary
┏━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃ path ┃ module ┃ inputs ┃ outputs ┃ flops ┃ vjp_flops ┃ params ┃
┡━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│ │ Foo │ float32[16,9] │ float32[16,2] │ 0 │ 0 │ │
├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼──────────────────────┤
│ Dense_0 │ Dense │ float32[16,9] │ float32[16,4] │ 0 │ 0 │ bias: float32[4] │
│ │ │ │ │ │ │ kernel: float32[9,4] │
│ │ │ │ │ │ │ │
│ │ │ │ │ │ │ 40 (160 B) │
├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼──────────────────────┤
│ Dense_1 │ Dense │ float32[16,4] │ float32[16,2] │ 0 │ 0 │ bias: float32[2] │
│ │ │ │ │ │ │ kernel: float32[4,2] │
│ │ │ │ │ │ │ │
│ │ │ │ │ │ │ 10 (40 B) │
├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼──────────────────────┤
│ │ │ │ │ │ Total │ 50 (200 B) │
└─────────┴────────┴───────────────┴───────────────┴───────┴───────────┴──────────────────────┘
```
Contributor guide
Assessment
This issue has not been assessed yet.