Diff Coverage

    tensors' storage for the lifetime of training.
    """
    if isinstance(x, torch.Tensor):
        return {"shape": tuple(x.shape), "dtype": str(x.dtype)}
    if isinstance(x, tuple):
        return tuple(_input_meta(v) for v in x)
    if isinstance(x, list):
        return [_input_meta(v) for v in x]
    if isinstance(x, dict):
        return {k: _input_meta(v) for k, v in x.items()}
    return x


def trace_model_graph(
    model: torch.nn.Module,

        trainable = self._trainable_params_in_state_order(
            self._joint_graph.state_fqns, state_is_param, fqn_to_param
        )
        if len(trainable) != len(grads):
            raise ValueError(
                f"Gradient count ({len(grads)}) does not match trainable "
                f"parameter count ({len(trainable)}). The traced graph and "
                f"the live model disagree on which parameters are trainable; "
                f"refusing to assign gradients to avoid silent misalignment."

        """
        trainable: List[torch.nn.Parameter] = []
        for idx, fqn in enumerate(state_fqns):
            if state_is_param is not None and not state_is_param[idx]:
                continue
            param = fqn_to_param.get(fqn)
            if param is not None and param.requires_grad:
                trainable.append(param)
        return trainable