def check_layer(
    layer: pz.Layer,
    argument: Any | None = None,
    initialize: bool = True,
) -> Any:
  """Checks that a layer has been configured correctly by tracing it.

  This function runs the layer under `jax.eval_shape` and passes in a dummy
  input that matches the expected structure, with arbitrary sizes for
  unspecified dimension variables. It then checks that the output shape matches
  the expected output shape.

  Note that not all layers may be able to fully encode their preconditions
  in their ``input_structure``; in this case it may be necessary to provide a
  dummy argument that matches the expected input structure.

  Args:
    layer: The layer to check.
    argument: An optional argument to pass to the layer. If not provided, a
      dummy argument will be created based on the layer's input structure. Can
      contain jax.ShapeDtypeStruct leaves.
    initialize: Whether to initialize any uninitialized parameters in the model.

  Returns:
    The traced output structure from the call, as a PyTree of
    jax.ShapeDtypeStruct leaves.

  Raises:
    ValueError: If the layer's input structure contains ANY but no dummy
      argument was provided.

    Note that the layer may also raise other exceptions if misconfigured.
  """

  def go(layer, argument):
    if initialize:
      layer = pz.nn.initialize_parameters(layer, jax.random.key(0))

    if argument is None:
      input_structure = layer.input_structure()
      if (
          not pz.select(input_structure)
          .at_instances_of(pz.chk.Wildcard)
          .is_empty()
      ):
        raise ValueError(
            "Cannot synthesize an input for a layer whose input structure"
            " contains ANY. Please provide an explicit dummy argument."
        )
      dimension_vars = pz.chk.get_dimension_variables(input_structure)
      new_subst = pz.chk.DimensionVariableSubstitution({}, {}, {})
      primes = _generate_primes()
      for _ in range(3):
        # Skip the small primes
        next(primes)
      for var in dimension_vars.size_variables:
        new_subst.size_variables[var] = next(primes)
      for var in dimension_vars.sequence_variables:
        new_subst.sequence_variables[var] = (next(primes), next(primes))
      for var in dimension_vars.mapping_variables:
        n1 = next(primes)
        n2 = next(primes)
        new_subst.mapping_variables[var] = {
            f"__trace_check_{n1}__": n1,
            f"__trace_check_{n2}__": n2,
        }
      argument_structure = pz.chk.full_substitute_dimension_variables(
          input_structure, new_subst
      )
      argument = (
          pz.select(argument_structure)
          .at_instances_of(pz.chk.ArraySpec)
          .apply(_generate_with_sensible_dtype)
      )

    matches = pz.chk.check_structure(argument, layer.input_structure())
    result = layer(argument)
    pz.chk.check_structure(result, layer.output_structure(), known_vars=matches)
    return result

  return jax.eval_shape(go, layer, argument)
Checks that a layer has been configured correctly by tracing it.

  This function runs the layer under `jax.eval_shape` and passes in a dummy
  input that matches the expected structure, with arbitrary sizes for
  unspecified dimension variables. It then checks that the output shape matches
  the expected output shape.

  Note that not all layers may be able to fully encode their preconditions
  in their ``input_structure``; in this case it may be necessary to provide a
  dummy argument that matches the expected input structure.

  Args:
    layer: The layer to check.
    argument: An optional argument to pass to the layer. If not provided, a
      dummy argument will be created based on the layer's input structure. Can
      contain jax.ShapeDtypeStruct leaves.
    initialize: Whether to initialize any uninitialized parameters in the model.

  Returns:
    The traced output structure from the call, as a PyTree of
    jax.ShapeDtypeStruct leaves.

  Raises:
    ValueError: If the layer's input structure contains ANY but no dummy
      argument was provided.

    Note that the layer may also raise other exceptions if misconfigured.
  
penzai/deprecated/v1/toolshed/check_layers_by_tracing.py