Incorrect flop count estimate when using flax.linen.remat_scan

When using linen's remat_scan functionality for constant-time compilation w.r.t. depth, our team is encountering an error in the estimated FLOP count for the compiled training step.

In particular, we are observing that the estimated flop count is _lower_ than the product of the batch size per device times the model parameters, which is impossible, since the gradient of the loss must be computed for each example separately, then summed.

The issue goes away when using a for-loop to process layers instead, even though the parameter count is the same for both models. The flop count estimate also changes slightly for jit vs pmap, but the issue of a clearly erroneous flop count is only observed when using flax.linen.remat_scan.

### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): MacOS latest
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: ```flax==0.6.11, jax=0.4.9, jaxlib==0.4.9```
- Python version: 3.10
- GPU/TPU model and memory: N/A
- CUDA version (if applicable): N/A

### Problem you have encountered:

### What you expected to happen:

### Logs, error messages, etc:

```
$ python3 scripts/flops_issue.py --use_remat_scan=False
I1116 17:09:23.125371 4452404672 xla_bridge.py:450] Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I1116 17:09:23.125525 4452404672 xla_bridge.py:450] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I1116 17:09:23.125777 4452404672 xla_bridge.py:450] Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'
I1116 17:09:23.125882 4452404672 xla_bridge.py:450] Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.
I1116 17:09:27.038714 4452404672 flops_issue.py:132] FLOP count estimate: 1091282.0
I1116 17:09:27.039010 4452404672 flops_issue.py:134] Param count: 20400
I1116 17:09:27.039087 4452404672 flops_issue.py:136] Example count: 8.0
I1116 17:09:27.039138 4452404672 flops_issue.py:138] Reasonable flop count min: 163200.0
I1116 17:09:27.039187 4452404672 flops_issue.py:140] Is reasonable estimate: True
I1116 17:09:29.560203 4452404672 flops_issue.py:132] FLOP count estimate: 1070881.0
I1116 17:09:29.560420 4452404672 flops_issue.py:134] Param count: 20400
I1116 17:09:29.560477 4452404672 flops_issue.py:136] Example count: 8.0
I1116 17:09:29.560518 4452404672 flops_issue.py:138] Reasonable flop count min: 163200.0
I1116 17:09:29.560559 4452404672 flops_issue.py:140] Is reasonable estimate: True
```

```
$ python3 scripts/flops_issue.py --use_remat_scan=True
I1116 17:10:44.239187 4791418304 xla_bridge.py:450] Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I1116 17:10:44.239346 4791418304 xla_bridge.py:450] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I1116 17:10:44.239608 4791418304 xla_bridge.py:450] Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'
I1116 17:10:44.239680 4791418304 xla_bridge.py:450] Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.
I1116 17:10:46.316076 4791418304 flops_issue.py:132] FLOP count estimate: 112855.0
I1116 17:10:46.316204 4791418304 flops_issue.py:134] Param count: 20400
I1116 17:10:46.316262 4791418304 flops_issue.py:136] Example count: 8.0
I1116 17:10:46.316305 4791418304 flops_issue.py:138] Reasonable flop count min: 163200.0
I1116 17:10:46.316347 4791418304 flops_issue.py:140] Is reasonable estimate: False
I1116 17:10:47.884852 4791418304 flops_issue.py:132] FLOP count estimate: 92454.0
I1116 17:10:47.884963 4791418304 flops_issue.py:134] Param count: 20400
I1116 17:10:47.885011 4791418304 flops_issue.py:136] Example count: 8.0
I1116 17:10:47.885046 4791418304 flops_issue.py:138] Reasonable flop count min: 163200.0
I1116 17:10:47.885081 4791418304 flops_issue.py:140] Is reasonable estimate: False
```

### Steps to reproduce:
Whenever possible, please provide a *minimal example*. Please consider submitting it as a Colab link.

```
from typing import Dict
from typing import Tuple
from typing import Type

import flax.linen as nn
import jax
import jax.numpy as jnp
import optax
from absl import app
from absl import flags
from absl import logging
from flax import jax_utils
from flax.training import common_utils
from flax.training.train_state import TrainState

FLAGS = flags.FLAGS
flags.DEFINE_boolean("use_remat_scan", True, "Use remat_scan?")

BSZ = 8
D_MODEL = 10
D_FF = 40
N_LAYER = 24

class MLP(nn.Module):
d_model: int
d_ff: int

@nn.compact
def __call__(self, x):
af1 = nn.Dense(self.d_ff)(x)
af2 = nn.Dense(self.d_model)(jax.nn.silu(af1))
return x + af2

class StackedMLPsRematted(nn.Module):
d_model: int
d_ff: int
n_layer: int

@nn.compact
def __call__(self, x):
return nn.remat_scan(MLP, lengths=(self.n_layer, 1))(
d_model=self.d_model,
d_ff=self.d_ff,
name="stack",
)(x)

class StackedMLPsListed(nn.Module):
d_model: int
d_ff: int
n_layer: int

@nn.compact
def __call__(self, x):
for _ in range(self.n_layer):
x = MLP(d_model=self.d_model, d_ff=self.d_ff)(x)
return x

def loss_fn(params, batch, cls):
predictions = cls(D_MODEL, D_FF, N_LAYER).apply({"params": params}, batch["inputs"])
loss_terms = optax.l2_loss(targets=batch["targets"], predictions=predictions)
loss = jnp.mean(loss_terms)
return loss

def train_op(
train_state: TrainState,
batch: Dict[str, jax.Array],
is_pmapped: bool,
cls: Type[nn.Module],
) -> Tuple[TrainState, jax.Array]:
loss, grads = jax.value_and_grad(loss_fn)(
train_state.params,
batch=batch,
cls=cls,
)
if is_pmapped:
loss, grads = jax.lax.pmean([loss, grads], axis_name="devices")
return train_state.apply_gradients(grads=grads), loss

def get_deterministic_trainstate_and_batch(cls):
x = jax.random.normal(jax.random.PRNGKey(0), [BSZ, D_MODEL])
y = jax.random.normal(jax.random.PRNGKey(1), [BSZ, D_MODEL])
batch = dict(inputs=x, targets=y)
ps = cls(D_MODEL, D_FF, N_LAYER).init({"params": jax.random.PRNGKey(2)}, x)
train_state = TrainState.create(
apply_fn=None,
params=ps["params"].unfreeze(),
tx=optax.sgd(learning_rate=0.01),
)
return train_state, batch

def pmapped_train_op_cost_analysis(cls):
train_state, batch = get_deterministic_trainstate_and_batch(cls)
p_train_op = jax.pmap(
train_op,
axis_name="devices",
donate_argnums=(0,),
static_broadcasted_argnums=(2, 3),
)
compiled = p_train_op.lower(
jax_utils.replicate(train_state),
common_utils.shard(batch),
True,
cls,
).compile()
cost_analysis = compiled.cost_analysis()
return cost_analysis

def jitted_train_op_cost_analysis(cls):
train_state, batch = get_deterministic_trainstate_and_batch(cls)
j_train_op = jax.jit(train_op, donate_argnums=(0,), static_argnums=(2, 3))
compiled = j_train_op.lower(
train_state,
batch,
False,
cls,
).compile()
cost_analysis = compiled.cost_analysis()
return cost_analysis

def evaluate_cost_analysis(train_state, cost_analysis):
if cost_analysis is not None:
n_flop = cost_analysis[0]["flops"]
logging.info(f"FLOP count estimate: {n_flop}")
n_param = sum(x.size for x in jax.tree_util.tree_leaves(train_state.params))
logging.info(f"Param count: {n_param}")
n_example = BSZ / jax.local_device_count()
logging.info(f"Example count: {n_example}")
n_min_flop_reasonable = n_example * n_param
logging.info(f"Reasonable flop count min: {n_min_flop_reasonable}")
is_reasonable = n_flop > n_min_flop_reasonable
logging.info(f"Is reasonable estimate: {is_reasonable}")

def main(argv):
del argv
cls = StackedMLPsRematted if FLAGS.use_remat_scan else StackedMLPsListed

cost_for_pmapped_train_op = pmapped_train_op_cost_analysis(cls)
train_state, _ = get_deterministic_trainstate_and_batch(cls)
evaluate_cost_analysis(train_state, cost_for_pmapped_train_op)

cost_for_jitted_train_op = jitted_train_op_cost_analysis(cls)
train_state, _ = get_deterministic_trainstate_and_batch(cls)
evaluate_cost_analysis(train_state, cost_for_jitted_train_op)

if __name__ == "__main__":
jax.config.config_with_absl()
app.run(main)
```