code: def _maybe_unknown(x: Any) -> pe.PartialVal:
      if isinstance(x, jax.ShapeDtypeStruct):
          return pe.PartialVal.unknown(
              core.ShapedArray(x.shape, x.dtype)
          )
      else:
          return pe.PartialVal.known(x)
  
file_path: flax/flax/core/partial_eval.py
imports: import functools
  from typing import Any
  import jax
  from jax import core
  from jax.extend import linear_util
  from jax.interpreters import partial_eval
  from flax import errors
label_desc: Use image node_id 1 for calling a global function with example usage: _maybe_unknown(x) and returns: pe, pe