# Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.
#
# 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
#
#     http://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.
# ==============================================================================
"""Transformation wrappers."""

from optax.transforms import _accumulation
from optax.transforms import _conditionality
from optax.transforms import _layouts
from optax.transforms import _masking


apply_if_finite = _conditionality.apply_if_finite
ApplyIfFiniteState = _conditionality.ApplyIfFiniteState
ConditionFn = _conditionality.ConditionFn
conditionally_mask = _conditionality.conditionally_mask
conditionally_transform = _conditionality.conditionally_transform
ConditionallyMaskState = _conditionality.ConditionallyMaskState
ConditionallyTransformState = _conditionality.ConditionallyTransformState
flatten = _layouts.flatten
masked = _masking.masked
MaskedNode = _masking.MaskedNode
MaskedState = _masking.MaskedState
MultiSteps = _accumulation.MultiSteps
MultiStepsState = _accumulation.MultiStepsState
ShouldSkipUpdateFunction = _accumulation.ShouldSkipUpdateFunction
skip_not_finite = _accumulation.skip_not_finite
skip_large_updates = _accumulation.skip_large_updates
