37 lines
2.0 KiB
Python
37 lines
2.0 KiB
Python
|
# Copyright 2022 The JAX Authors.
|
||
|
|
#
|
||
|
|
# 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, cumlogsumexp,
|
||
|
|
cumlogsumexp_p, cumprod,
|
||
|
|
cumprod_p, cumsum, cumsum_p,
|
||
|
|
cumred_reduce_window_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.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)
|
||
|
|
# TODO(mattjj): fix dependent library which expects optimization_barrier_p here
|
||
|
|
from jax._src.ad_checkpoint import optimization_barrier_p
|