# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Module for the control flow primitives."""
from jax._src.lax.control_flow.loops import (associative_scan, cummax, cummax_p,
                                             cummin, cummin_p, cumprod,
                                             cumprod_p, cumsum, cumsum_p,
                                             cumred_tpu_impl, fori_loop, map,
                                             scan, scan_bind, scan_p,
                                             _scan_impl, while_loop, while_p)
from jax._src.lax.control_flow.conditionals import cond, cond_p, switch
from jax._src.lax.control_flow.remat_impl import (remat_impl,
                                                  optimization_barrier_p)
from jax._src.lax.control_flow.solves import (custom_linear_solve, custom_root,
                                              _custom_linear_solve_impl,
                                              linear_solve_p)

from jax._src.lax.control_flow.common import allowed_effects
# Private utilities used elsewhere in JAX
# TODO(sharadmv): lift them into a more common place
from jax._src.lax.control_flow.common import (_initial_style_open_jaxpr,
                                              _initial_style_jaxpr,
                                              _initial_style_jaxprs_with_common_consts,
                                              _check_tree_and_avals)