mirror of
https://github.com/ROCm/jax.git
synced 2025-04-25 10:26:07 +00:00
d51b8e6839
This change re-introduces symbolic zero support for `custom_vjp`. This time: * The forward rule API is slightly different, accepting two-field records at pytree leaves rather than pairs. * In the default setting where symbolic_zeros is not set, there are no new requirements from pytree node definitions that are involved in the primal arguments. This avoids any change in behavior on the default path. In particular, custom pytree node definitions that aren't completely polymorphic in unflattening can remain as is. * There is an additional test involving a custom pytree node.
38 lines
1.3 KiB
Python
38 lines
1.3 KiB
Python
# Copyright 2020 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.
|
|
|
|
# Note: import <name> as <name> is required for names to be exported.
|
|
# See PEP 484 & https://github.com/google/jax/issues/7570
|
|
|
|
from jax._src.custom_derivatives import (
|
|
_initial_style_jaxpr,
|
|
_sum_tangents,
|
|
_zeros_like_pytree,
|
|
closure_convert as closure_convert,
|
|
custom_gradient as custom_gradient,
|
|
custom_jvp as custom_jvp,
|
|
custom_jvp_call_p as custom_jvp_call_p,
|
|
custom_jvp_call_jaxpr_p as custom_jvp_call_jaxpr_p,
|
|
custom_vjp as custom_vjp,
|
|
custom_vjp_call_p as custom_vjp_call_p,
|
|
custom_vjp_call_jaxpr_p as custom_vjp_call_jaxpr_p,
|
|
custom_vjp_primal_tree_values as custom_vjp_primal_tree_values,
|
|
CustomVJPPrimal as CustomVJPPrimal,
|
|
linear_call as linear_call,
|
|
)
|
|
|
|
from jax._src.ad_util import (
|
|
SymbolicZero as SymbolicZero
|
|
)
|