Explicitly export "struct" from module "flax".

This allows the following:
```
import flax

@flax.struct.dataclass
class C: ...
```

The current usages of the above pattern in flax only work
accidentally because there already are imports for "flax.struct"
in other files.

See also https://docs.python.org/3/reference/simple_stmts.html#the-import-statement

PiperOrigin-RevId: 552452298

flax/__init__.py

@@ -24,6 +24,7 @@
 from . import jax_utils
 from . import linen
 from . import serialization
+from . import struct
 from . import traverse_util
 
 # DO NOT REMOVE - Marker for internal deprecated API.