text: # Copyright 2018 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.
  
  # Set default logging level before any logging happens.
  import os as _os
  _os.environ.setdefault('TF_CPP_MIN_LOG_LEVEL', '1')
  del _os
  
  # flake8: noqa: F401
  from .config import config
  from .api import (
    ad,  # TODO(phawkins): update users to avoid this.
    argnums_partial,  # TODO(phawkins): update Haiku to not use this.
    checkpoint,
    curry,  # TODO(phawkins): update users to avoid this.
    custom_ivjp,
    custom_gradient,
    custom_jvp,
    custom_vjp,
    custom_transforms,
    defjvp,
    defjvp_all,
    defvjp,
    defvjp_all,
    device_count,
    device_get,
    device_put,
    devices,
    disable_jit,
    eval_shape,
    flatten_fun_nokwargs,  # TODO(phawkins): update users to avoid this.
    float0,
    grad,
    hessian,
    host_count,
    host_id,
    host_ids,
    invertible,
    jacobian,
    jacfwd,
    jacrev,
    jit,
    jvp,
    local_device_count,
    local_devices,
    linearize,
    linear_transpose,
    make_jaxpr,
    mask,
    named_call,
    partial,  # TODO(phawkins): update callers to use functools.partial.
    pmap,
    pxla,  # TODO(phawkins): update users to avoid this.
    remat,
    shapecheck,
    ShapedArray,
    ShapeDtypeStruct,
    soft_pmap,
    # TODO(phawkins): hide tree* functions from jax, update callers to use
    # jax.tree_util.
    treedef_is_leaf,
    tree_flatten,
    tree_leaves,
    tree_map,
    tree_multimap,
    tree_structure,
    tree_transpose,
    tree_unflatten,
    value_and_grad,
    vjp,
    vmap,
    xla,  # TODO(phawkins): update users to avoid this.
    xla_computation,
  )
  from .version import __version__
  
  # These submodules are separate because they are in an import cycle with
  # jax and rely on the names imported above.
  from . import image
  from . import lax
  from . import nn
  from . import profiler
  from . import random
  
  def _init():
    from . import numpy # side-effecting import sets up operator overloads
  
  _init()
  del _init
  