Bug - Cannot access attribute "replace" for class "c123"

Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.

### System information
- OS Platform and Distribution (WSL 2 , Ubuntu 22):
- jax - 0.4.30
- flax - 0.10.0
- jaxlib - 0.4.29
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`:
- Python version: 3.10.12
- GPU/TPU model and memory: CPU, 8gb RAM
- CUDA version (if applicable): -

### Problem you have encountered:
Replace function in flax is shown as an error, as if it does not exist.

### What you expected to happen:
The GUI shows no errors. It's quite annoying given that .replace is something to be called often.

### Logs, error messages, etc:
Cannot access attribute "replace" for class "c123"
Attribute "replace" is unknownPylance[reportAttributeAccessIssue](https://github.com/microsoft/pyright/blob/main/docs/configuration.md#reportAttributeAccessIssue)

### Steps to reproduce:
Whenever possible, please provide a *minimal example*.

```
@dataclass
class c123:
variable1: int = 0

c123instance = c123(variable1=0).replace(variable1=1)

print(c123instance)
>>> c123(variable1=1)
```

---
Also, I would ask, how is `.replace` function working if it calls `self` in it's implementation? To my knowledge and to Jax docs, `self` should throw an error every time there is jit.