code: def _foldl_jax(
      fn,
      elems,
      initializer=None,
      parallel_iterations=10,  # pylint: disable=unused-argument
      back_prop=True,
      swap_memory=False,
      name=None,
  ):  # pylint: disable=unused-argument
      """tf.foldl, in JAX."""
      if initializer is None:
          initializer = nest.map_structure(lambda el: el[0], elems)
          elems = nest.map_structure(lambda el: el[1:], elems)
      if len(set(nest.flatten(nest.map_structure(len, elems)))) != 1:
          raise ValueError(
              "Mismatched element sizes: {}".format(
                  nest.map_structure(len, elems)
              )
          )
      from jax import lax  # pylint: disable=g-import-not-at-top
  
      return lax.scan(
          lambda carry, el: (fn(carry, el), None), initializer, elems
      )[0]
  
file_path: probability/tensorflow_probability/python/internal/backend/numpy/functional_ops.py
imports: import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import nest
label_desc: Use image node_id 1 for calling a global function with example usage: _foldl_jax(fn, elems, initializer, parallel_iterations, back_prop, swap_memory, name) and returns: lax