code: def get_params_batch_length(params: PyTree) -> int:
      """Get the batch length from a params dictionary."""
      return jax.tree_util.tree_leaves(params)[0].shape[0]
  
imports: import functools
  from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
  import bayeux
  import jax
  import jax.numpy
  from jaxtyping import PyTree
  import matplotlib.pyplot
  import numpy
  import pandas
  from tensorflow_probability.python.experimental.autobnn import bnn
  from tensorflow_probability.python.experimental.autobnn import util
  from tensorflow_probability.python.experimental.timeseries import metrics
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 11 for calling a global function with example usage: get_params_batch_length(params) and returns: jax