forked from google/jax
-
Notifications
You must be signed in to change notification settings - Fork 2
/
linear_util.py
28 lines (25 loc) · 1.02 KB
/
linear_util.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
# Copyright 2018 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
# TODO(jakevdp): deprecate these and remove this module.
from jax._src.linear_util import (
StoreException as StoreException,
WrappedFun as WrappedFun,
cache as cache,
merge_linear_aux as merge_linear_aux,
transformation as transformation,
transformation_with_aux as transformation_with_aux,
wrap_init as wrap_init,
)