# Copyright 2022 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.
# ==============================================================================
"""Stubs for jax.*

NOTE: This file is generated from templates/jax.pyi.

To regenerate, run the following from the tensor_annotations directory:
   tools/render_jax_library_template.py
"""

from typing import Any

_sentinel: int


ShapeDtypeStruct: Any

ShapedArray: Any

__builtins__: Any

__cached__: Any

__doc__: Any

__file__: Any

__loader__: Any

__name__: Any

__package__: Any

__path__: Any

__spec__: Any

__version__: Any

__version_info__: Any

abstract_arrays: Any

ad: Any

ad_checkpoint: Any

api_util: Any

block_until_ready: Any

check_tracer_leaks: Any

checking_leaks: Any

checkpoint: Any

checkpoint_policies: Any

clear_backends: Any

closure_convert: Any

config: Any

core: Any

curry: Any

custom_batching: Any

custom_derivatives: Any

custom_gradient: Any

custom_jvp: Any

custom_transpose: Any

custom_vjp: Any

debug: Any

debug_infs: Any

debug_nans: Any

default_backend: Any

default_device: Any

default_matmul_precision: Any

default_prng_impl: Any

device_count: Any

device_get: Any

device_put: Any

device_put_replicated: Any

device_put_sharded: Any

devices: Any

disable_jit: Any

distributed: Any

dtypes: Any

effects_barrier: Any

enable_checks: Any

enable_custom_prng: Any

enable_custom_vjp_by_custom_transpose: Any

ensure_compile_time_eval: Any

errors: Any

eval_shape: Any

experimental: Any

flatten_fun_nokwargs: Any

float0: Any

grad: Any

hessian: Any

host_count: Any

host_id: Any

host_ids: Any

image: Any

interpreters: Any

jacfwd: Any

jacobian: Any

jacrev: Any

jax: Any

jax2tf_associative_scan_reductions: Any

jax_google: Any

jaxlib: Any

jit: Any

jvp: Any

lax: Any

lib: Any

linear_transpose: Any

linear_util: Any

linearize: Any

local_device_count: Any

local_devices: Any

log_compiles: Any

make_jaxpr: Any

named_call: Any

named_scope: Any

nn: Any

numpy_dtype_promotion: Any

numpy_rank_promotion: Any

ops: Any

pmap: Any

print_environment_info: Any

process_count: Any

process_index: Any

profiler: Any

pure_callback: Any

pxla: Any

random: Any

remat: Any

scipy: Any

soft_pmap: Any

stages: Any

transfer_guard: Any

transfer_guard_device_to_device: Any

transfer_guard_device_to_host: Any

transfer_guard_host_to_device: Any

tree_flatten: Any

tree_leaves: Any

tree_map: Any

tree_structure: Any

tree_transpose: Any

tree_unflatten: Any

tree_util: Any

treedef_is_leaf: Any

util: Any

value_and_grad: Any

version: Any

vjp: Any

vmap: Any

xla: Any

xla_computation: Any
