def _generate_with_sensible_dtype(structure: pz.chk.ArraySpec):
  """Constructs a zeros array with a reasonable dtype and a given structure."""
  if structure.dtype == jax.dtypes.prng_key:
    new_dtype = jax.eval_shape(jax.random.key, 0).dtype
  else:
    concrete_default_dtypes = [np.float32, np.int32, np.bool_, np.complex64]
    for dtype in concrete_default_dtypes:
      if jnp.issubdtype(dtype, structure.dtype):
        new_dtype = dtype
        break
    else:
      new_dtype = jax.dtypes.result_type(structure.dtype)
  adj_structure = pz.chk.ArraySpec(
      shape=structure.shape, named_shape=structure.named_shape, dtype=new_dtype
  ).into_pytree()
  return jax.tree_util.tree_map(jnp.zeros_like, adj_structure)
Constructs a zeros array with a reasonable dtype and a given structure.
penzai/deprecated/v1/toolshed/check_layers_by_tracing.py