code: def fork(self, pattern: Pattern) -> "RngStream":
      if pattern is None:
          # broadcast key
          key = self.key
          count_path = [*self.counts, 0]
          self.counts[-1] += 1
      else:
          key = self.make_rng()
          # split key
          if isinstance(pattern, int):
              key = jax.random.split(key, pattern)
          else:
              num_splits = int(
                  np.prod([x for x in pattern if x is not None])
              )
              axis_size = tuple(
                  x if x is not None else 1 for x in pattern
              )
              # reshape key
              key = jax.random.split(key, num_splits).reshape(
                  *axis_size
              )
          count_path = [0]
      return RngStream(key, count_path)
  
imports: from __future__ import annotations
  import dataclasses
  import hashlib
  import typing
  import jax
  import numpy
  from flax.experimental.nnx.nnx import errors, filterlib, tracers
label_desc: Use image node_id 2 for calling the RngStream obj's underlying member method code with example usage: obj.fork(pattern) and returns: RngStream