nnx's GPU memory usage is significantly larger than Pytorch

I'm trying to build a few billion-parameter scale Transformer model, and trying both nnx and pytorch to compare the pros and cons before choosing the library.

I implemented the same model with both libraries, and found out that nnx allocates significantly more GPU memory than Pytorch. (I set `XLA_PYTHON_CLIENT_PREALLOCATE=false` in nnx)

At first, I thought that this is caused by the different attention kernels (flash attention kernel in Pytorch vs vanilla attention in nnx.dot_product_attention/jax.nn.dot_product_attention).

However, I could see the same pattern in the following simple test scripts:

```py
import os
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
import time

import jax
from flax import nnx
from jax import numpy as jnp

B = 128
D = 2048

x = jnp.ones((B, D),dtype=jnp.float32)
# Total memory usage after x : 557Mb

m = nnx.Linear(D, D, rngs=nnx.Rngs(0))
# Total memory usage after m : 593Mb

def run(m, q):
return m(q)

runjit = nnx.jit(run)

while True:
tt = time.time()
y = runjit(m, x)
jax.block_until_ready(y)
# total memory after y : 755Mb
print(time.time() - tt)
```

```py
import time
import torch

B = 128
D = 2048

x = torch.ones(B,D).to("cuda:0")
# Total memory usage after x : 531Mb

run = torch.nn.Linear(D, D).to("cuda:0")
# Total memory usage after m : 547Mb

runjit = torch.compile(run)

while True:
tt = time.time()
y = runjit(x)
torch.cuda.synchronize()
# total memory after y : 653Mb (with or without torch.compile)
print(time.time() - tt)
```

The gap became larger when I tested with larger models:

```py
import os
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
import time

import jax
from flax import nnx
from jax import numpy as jnp

B = 128
D = 2048

x = jnp.ones((B, D),dtype=jnp.float32)
# Total memory usage after x : 557Mb

class Model(nnx.Module):
def __init__(self,rngs):
self.layers = [nnx.Linear(D,D,rngs=nnx.Rngs(0)) for _ in range(18)]
def __call__(self,x):
for layer in self.layers:
x = nnx.relu(layer(x))
return x

m = Model(rngs=nnx.Rngs(0))
# Total memory usage after m : 1073Mb

def run(m, q):
return m(q)

runjit = nnx.jit(run)

while True:
tt = time.time()
y = runjit(m, x)
jax.block_until_ready(y)
# total memory after y : 1651Mb
print(time.time() - tt)
```
```py
import time
import torch
from torch import nn
from torch.nn import functional as F

B = 128
D = 2048

x = torch.ones(B,D).to("cuda:0")
# Total memory usage after x : 531Mb

class Model(nn.Module):
def __init__(self):
super().__init__()
self.layers = nn.ModuleList()
for _ in range(18):
self.layers.append(nn.Linear(D,D))

def forward(self,x):
for layer in self.layers:
x = F.relu(layer(x))
return x

run = Model().to("cuda:0")
# Total memory usage after m : 819Mb

runjit = torch.compile(run)

while True:
tt = time.time()
y = runjit(x)
torch.cuda.synchronize()
# total memory after y : 1043Mb (with torch.compile), 1045Mb (without torch.compile)
print(time.time() - tt)
```

In the 3 billion transformer, this gap becomes dozens of Gbs, so I had to use smaller batch size for nnx model.

Is it normal? or is there a trick to optimize the memory in nnx?