code: def _get_params_shapes_in_lstm(
      input_size: int,
      hidden_size: int,
      num_layers: int,
      bidirectional: bool,
  ) -> list[Shape]:
      """Get flat param shapes in LSTM. See module docstring for layout."""
      layer_shapes = []
      num_directions = 2 if bidirectional else 1
      num_pseudo_layers = num_layers * num_directions
      linear_weights = [_W_ih_l, _W_hh_l]
      for i in range(num_pseudo_layers):
          for w_kind in linear_weights:
              layer_shape = w_kind(
                  i, input_size, hidden_size, bidirectional
              )
              layer_shapes.append(layer_shape)
  
      bias_weights = [_b_ih_l, _b_hh_l]
      for i in range(num_pseudo_layers):
          for w_kind in bias_weights:
              layer_shape = w_kind(
                  i, input_size, hidden_size, bidirectional
              )
              layer_shapes.append(layer_shape)
      return layer_shapes
  
imports: from functools import partial
  import math
  from typing import Any
  import jax
  import numpy
  from jax._src import core
  from jax.interpreters import mlir
  from jax.interpreters import xla
  from jax._src.custom_derivatives import custom_vjp
  from jax._src.typing import Array, Shape
  from jax._src.lax import lax
  import jax.numpy
label_desc: Use image node_id 5 for calling a global function with example usage: _get_params_shapes_in_lstm(input_size, hidden_size, num_layers, bidirectional) and returns: layer_shapes