code: def make_extra_packages_flax():
      # Packages needed for the flax example.
      # Required for the experimental tfx/examples using Flax, e.g.,
      # tfx/examples/penguin.
      return [
          "jax<1",
          "jaxlib<1",
          "flax<1",
          "optax<1",
      ]
  
file_path: tfx/tfx/dependencies.py
imports: import os
label_desc: Use image node_id 12 for calling a global function with example usage: make_extra_packages_flax() and returns: list