code: def flatten(obs: TensorType, framework: str) -> TensorType:
      """Flatten the given tensor."""
      if framework in ["tf2", "tf"]:
          return tf1.keras.layers.Flatten()(obs)
      elif framework == "torch":
          assert torch is not None
          return torch.flatten(obs, start_dim=1)
      else:
          raise NotImplementedError("flatten", framework)
  
imports: from collections import OrderedDict
  import contextlib
  import gymnasium
  from gymnasium.spaces import Space
  import numpy
  from typing import Dict, List, Any, Union
  from ray.rllib.models.preprocessors import get_preprocessor, RepeatedValuesPreprocessor
  from ray.rllib.models.repeated_values import RepeatedValues
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.view_requirement import ViewRequirement
  from ray.rllib.utils import NullContextManager
  from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI
  from ray.rllib.utils.deprecation import Deprecated
  from ray.rllib.utils.framework import try_import_tf, try_import_torch, TensorType
  from ray.rllib.utils.spaces.repeated import Repeated
  from ray.rllib.utils.typing import ModelConfigDict, ModelInputDict, TensorStructType
label_desc: Use image node_id 1 for calling a global function with example usage: flatten(obs, framework) and returns: tf1, torch