google / google/flax

LazyRNG can accidentally have key sharing across layers.

Open
#2,157 2 comments 0 reactions 1 assignee Claimed by @jheek View on GitHub
Priority: P1 - soon
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

User reported repo:
```
import jax
from flax import linen as nn

class Leaf(nn.Module):
def __call__(self, x):
return x + jax.random.randint(self.make_rng("rng"), (), 0, 100)

class Node(nn.Module):
leaf_name: str
@nn.compact
def __call__(self, x):
return Leaf(name=self.leaf_name)(x)

class Model(nn.Module):
@nn.compact
def __call__(self, x):
return (Node(name="ab", leaf_name="cdef")(x),
Node(name="abc", leaf_name="def")(x),
)

print(Model().apply({}, 0, rngs={"rng": jax.random.PRNGKey(33)})
# (DeviceArray(23, dtype=int32), DeviceArray(23, dtype=int32))
```

This occurs because we don't add a separator at:
https://github.com/google/flax/blob/main/flax/core/scope.py#L102

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.