code: def flatten(x: JaxArray) -> JaxArray:
      """Flattens input tensor to a 2D tensor.
  
      Args:
          x: input tensor with dimensions (n_1, n_2, ..., n_k)
  
      Returns:
          The input tensor reshaped to two dimensions (n_1, n_prod),
          where n_prod is equal to the product of n_2 to n_k.
      """
      return x.reshape([x.shape[0], -1])
  
file_path: objax/objax/functional/core/ops.py
imports: from typing import Union, Tuple
  import jax.nn
  from jax import numpy, lax
  from objax import util
  from objax.constants import Interpolate
  from objax.typing import JaxArray
label_desc: Use image node_id 1 for calling a global function with example usage: flatten(x) and returns: x