HamiltonZero

Model card Files Files and versions

xet

Community

HamiltonZero / src /kfac_jax /__init__.py

Twobombs

takeover359

Duplicate from simulacra-research/HamiltonZero

71a70fb 2 months ago

Raw History Blame Contribute Delete

7.77 kB

# 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.
"""kfac-jax public APIs."""

from kfac_jax._src import curvature_blocks
from kfac_jax._src import curvature_estimator
from kfac_jax._src import layers_and_loss_tags
from kfac_jax._src import loss_functions
from kfac_jax._src import optimizer
from kfac_jax._src import patches_second_moment
from kfac_jax._src import tag_graph_matcher
from kfac_jax._src import tracer
from kfac_jax._src import utils


__version__ = "0.0.8"

# Patches Second Moments
patches_moments = patches_second_moment.patches_moments
patches_moments_explicit = patches_second_moment.patches_moments_explicit

# Layers and loss tags
LayerData = layers_and_loss_tags.LayerData
LayerMetaData = layers_and_loss_tags.LayerMetaData
LossTag = layers_and_loss_tags.LossTag
LayerTag = layers_and_loss_tags.LayerTag
register_generic = layers_and_loss_tags.register_generic
register_dense = layers_and_loss_tags.register_dense
register_conv2d = layers_and_loss_tags.register_conv2d
register_scale_and_shift = layers_and_loss_tags.register_scale_and_shift

# Tag graph matcher
auto_register_tags = tag_graph_matcher.auto_register_tags

# Tracer
ProcessedJaxpr = tracer.ProcessedJaxpr
LayerVjpData = tracer.LayerVjpData
loss_tags_vjp = tracer.loss_tags_vjp
loss_tags_jvp = tracer.loss_tags_jvp
loss_tags_hvp = tracer.loss_tags_hvp
layer_tags_vjp = tracer.layer_tags_vjp

# Loss functions
LossFunction = loss_functions.LossFunction
NegativeLogProbLoss = loss_functions.NegativeLogProbLoss
DistributionNegativeLogProbLoss = loss_functions.DistributionNegativeLogProbLoss
NormalMeanNegativeLogProbLoss = loss_functions.NormalMeanNegativeLogProbLoss
NormalMeanVarianceNegativeLogProbLoss = (
    loss_functions.NormalMeanVarianceNegativeLogProbLoss)
MultiBernoulliNegativeLogProbLoss = (
    loss_functions.MultiBernoulliNegativeLogProbLoss)
CategoricalLogitsNegativeLogProbLoss = (
    loss_functions.CategoricalLogitsNegativeLogProbLoss)
OneHotCategoricalLogitsNegativeLogProbLoss = (
    loss_functions.OneHotCategoricalLogitsNegativeLogProbLoss)
register_sigmoid_cross_entropy_loss = (
    loss_functions.register_sigmoid_cross_entropy_loss)
register_multi_bernoulli_predictive_distribution = (
    loss_functions.register_multi_bernoulli_predictive_distribution)
register_softmax_cross_entropy_loss = (
    loss_functions.register_softmax_cross_entropy_loss)
register_categorical_predictive_distribution = (
    loss_functions.register_categorical_predictive_distribution)
register_squared_error_loss = loss_functions.register_squared_error_loss
register_normal_predictive_distribution = (
    loss_functions.register_normal_predictive_distribution)

# Curvature blocks
CurvatureBlock = curvature_blocks.CurvatureBlock
ScaledIdentity = curvature_blocks.ScaledIdentity
Diagonal = curvature_blocks.Diagonal
Full = curvature_blocks.Full
KroneckerFactored = curvature_blocks.KroneckerFactored
NaiveDiagonal = curvature_blocks.NaiveDiagonal
NaiveFull = curvature_blocks.NaiveFull
NaiveTNT = curvature_blocks.NaiveTNT
DenseDiagonal = curvature_blocks.DenseDiagonal
DenseFull = curvature_blocks.DenseFull
DenseTwoKroneckerFactored = curvature_blocks.DenseTwoKroneckerFactored
RepeatedDenseKroneckerFactored = curvature_blocks.RepeatedDenseKroneckerFactored
DenseTNT = curvature_blocks.DenseTNT
Conv2DDiagonal = curvature_blocks.Conv2DDiagonal
Conv2DFull = curvature_blocks.Conv2DFull
Conv2DTwoKroneckerFactored = curvature_blocks.Conv2DTwoKroneckerFactored
Conv2DTNT = curvature_blocks.Conv2DTNT
ScaleAndShiftDiagonal = curvature_blocks.ScaleAndShiftDiagonal
ScaleAndShiftFull = curvature_blocks.ScaleAndShiftFull
set_max_parallel_elements = curvature_blocks.set_max_parallel_elements
get_max_parallel_elements = curvature_blocks.get_max_parallel_elements
set_default_eigen_decomposition_threshold = (
    curvature_blocks.set_default_eigen_decomposition_threshold)
get_default_eigen_decomposition_threshold = (
    curvature_blocks.get_default_eigen_decomposition_threshold)

# Curvature estimators
CurvatureEstimator = curvature_estimator.CurvatureEstimator
BlockDiagonalCurvature = curvature_estimator.BlockDiagonalCurvature
ExplicitExactCurvature = curvature_estimator.ExplicitExactCurvature
ImplicitExactCurvature = curvature_estimator.ImplicitExactCurvature
set_default_tag_to_block_ctor = (
    curvature_estimator.set_default_tag_to_block_ctor)
get_default_tag_to_block_ctor = (
    curvature_estimator.get_default_tag_to_block_ctor)
OptaxPreconditioner = curvature_estimator.OptaxPreconditioner
OptaxPreconditionState = curvature_estimator.OptaxPreconditionState

# Optimizers
Optimizer = optimizer.Optimizer

HAIKU_BIASES = optimizer.HAIKU_BIASES
HAIKU_BIASES_AND_NORMS = optimizer.HAIKU_BIASES_AND_NORMS

__all__ = (
    # Modules
    "utils",
    "patches_second_moment",
    "layers_and_loss_tags",
    "loss_functions",
    "tag_graph_matcher",
    "tracer",
    "curvature_blocks",
    "curvature_estimator",
    "optimizer",
    # Patches second moments
    "patches_moments",
    "patches_moments_explicit",
    # Layer and loss tags
    "LossTag",
    "LayerTag",
    "register_generic",
    "register_dense",
    "register_conv2d",
    "register_scale_and_shift",
    # Tag graph matcher
    "auto_register_tags",
    # Tracer
    "ProcessedJaxpr",
    "loss_tags_vjp",
    "loss_tags_jvp",
    "loss_tags_hvp",
    "layer_tags_vjp",
    # Loss functions
    "LossFunction",
    "NegativeLogProbLoss",
    "DistributionNegativeLogProbLoss",
    "NormalMeanNegativeLogProbLoss",
    "NormalMeanVarianceNegativeLogProbLoss",
    "MultiBernoulliNegativeLogProbLoss",
    "CategoricalLogitsNegativeLogProbLoss",
    "OneHotCategoricalLogitsNegativeLogProbLoss",
    "register_sigmoid_cross_entropy_loss",
    "register_multi_bernoulli_predictive_distribution",
    "register_softmax_cross_entropy_loss",
    "register_categorical_predictive_distribution",
    "register_squared_error_loss",
    "register_normal_predictive_distribution",
    # Curvature blocks
    "CurvatureBlock",
    "ScaledIdentity",
    "Diagonal",
    "Full",
    "KroneckerFactored",
    "NaiveDiagonal",
    "NaiveFull",
    "NaiveTNT",
    "DenseDiagonal",
    "DenseFull",
    "DenseTwoKroneckerFactored",
    "RepeatedDenseKroneckerFactored",
    "DenseTNT",
    "Conv2DDiagonal",
    "Conv2DFull",
    "Conv2DTwoKroneckerFactored",
    "Conv2DTNT",
    "ScaleAndShiftDiagonal",
    "ScaleAndShiftFull",
    "set_max_parallel_elements",
    "get_max_parallel_elements",
    "set_default_eigen_decomposition_threshold",
    "get_default_eigen_decomposition_threshold",
    # Estimators
    "CurvatureEstimator",
    "BlockDiagonalCurvature",
    "ExplicitExactCurvature",
    "ImplicitExactCurvature",
    "set_default_tag_to_block_ctor",
    "get_default_tag_to_block_ctor",
    # Optimizers
    "Optimizer",
)

#  _________________________________________
# / Please don't use symbols in `_src` they \
# \ are not part of the KFAC Jax public API./
#  -----------------------------------------
#         \   ^__^
#          \  (oo)\_______
#             (__)\       )\/\
#                 ||----w |
#                 ||     ||
#
try:
  del _src  # pylint: disable=undefined-variable
except NameError:
  pass

	# 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.
	"""kfac-jax public APIs."""

	from kfac_jax._src import curvature_blocks
	from kfac_jax._src import curvature_estimator
	from kfac_jax._src import layers_and_loss_tags
	from kfac_jax._src import loss_functions
	from kfac_jax._src import optimizer
	from kfac_jax._src import patches_second_moment
	from kfac_jax._src import tag_graph_matcher
	from kfac_jax._src import tracer
	from kfac_jax._src import utils


	__version__ = "0.0.8"

	# Patches Second Moments
	patches_moments = patches_second_moment.patches_moments
	patches_moments_explicit = patches_second_moment.patches_moments_explicit

	# Layers and loss tags
	LayerData = layers_and_loss_tags.LayerData
	LayerMetaData = layers_and_loss_tags.LayerMetaData
	LossTag = layers_and_loss_tags.LossTag
	LayerTag = layers_and_loss_tags.LayerTag
	register_generic = layers_and_loss_tags.register_generic
	register_dense = layers_and_loss_tags.register_dense
	register_conv2d = layers_and_loss_tags.register_conv2d
	register_scale_and_shift = layers_and_loss_tags.register_scale_and_shift

	# Tag graph matcher
	auto_register_tags = tag_graph_matcher.auto_register_tags

	# Tracer
	ProcessedJaxpr = tracer.ProcessedJaxpr
	LayerVjpData = tracer.LayerVjpData
	loss_tags_vjp = tracer.loss_tags_vjp
	loss_tags_jvp = tracer.loss_tags_jvp
	loss_tags_hvp = tracer.loss_tags_hvp
	layer_tags_vjp = tracer.layer_tags_vjp

	# Loss functions
	LossFunction = loss_functions.LossFunction
	NegativeLogProbLoss = loss_functions.NegativeLogProbLoss
	DistributionNegativeLogProbLoss = loss_functions.DistributionNegativeLogProbLoss
	NormalMeanNegativeLogProbLoss = loss_functions.NormalMeanNegativeLogProbLoss
	NormalMeanVarianceNegativeLogProbLoss = (
	loss_functions.NormalMeanVarianceNegativeLogProbLoss)
	MultiBernoulliNegativeLogProbLoss = (
	loss_functions.MultiBernoulliNegativeLogProbLoss)
	CategoricalLogitsNegativeLogProbLoss = (
	loss_functions.CategoricalLogitsNegativeLogProbLoss)
	OneHotCategoricalLogitsNegativeLogProbLoss = (
	loss_functions.OneHotCategoricalLogitsNegativeLogProbLoss)
	register_sigmoid_cross_entropy_loss = (
	loss_functions.register_sigmoid_cross_entropy_loss)
	register_multi_bernoulli_predictive_distribution = (
	loss_functions.register_multi_bernoulli_predictive_distribution)
	register_softmax_cross_entropy_loss = (
	loss_functions.register_softmax_cross_entropy_loss)
	register_categorical_predictive_distribution = (
	loss_functions.register_categorical_predictive_distribution)
	register_squared_error_loss = loss_functions.register_squared_error_loss
	register_normal_predictive_distribution = (
	loss_functions.register_normal_predictive_distribution)

	# Curvature blocks
	CurvatureBlock = curvature_blocks.CurvatureBlock
	ScaledIdentity = curvature_blocks.ScaledIdentity
	Diagonal = curvature_blocks.Diagonal
	Full = curvature_blocks.Full
	KroneckerFactored = curvature_blocks.KroneckerFactored
	NaiveDiagonal = curvature_blocks.NaiveDiagonal
	NaiveFull = curvature_blocks.NaiveFull
	NaiveTNT = curvature_blocks.NaiveTNT
	DenseDiagonal = curvature_blocks.DenseDiagonal
	DenseFull = curvature_blocks.DenseFull
	DenseTwoKroneckerFactored = curvature_blocks.DenseTwoKroneckerFactored
	RepeatedDenseKroneckerFactored = curvature_blocks.RepeatedDenseKroneckerFactored
	DenseTNT = curvature_blocks.DenseTNT
	Conv2DDiagonal = curvature_blocks.Conv2DDiagonal
	Conv2DFull = curvature_blocks.Conv2DFull
	Conv2DTwoKroneckerFactored = curvature_blocks.Conv2DTwoKroneckerFactored
	Conv2DTNT = curvature_blocks.Conv2DTNT
	ScaleAndShiftDiagonal = curvature_blocks.ScaleAndShiftDiagonal
	ScaleAndShiftFull = curvature_blocks.ScaleAndShiftFull
	set_max_parallel_elements = curvature_blocks.set_max_parallel_elements
	get_max_parallel_elements = curvature_blocks.get_max_parallel_elements
	set_default_eigen_decomposition_threshold = (
	curvature_blocks.set_default_eigen_decomposition_threshold)
	get_default_eigen_decomposition_threshold = (
	curvature_blocks.get_default_eigen_decomposition_threshold)

	# Curvature estimators
	CurvatureEstimator = curvature_estimator.CurvatureEstimator
	BlockDiagonalCurvature = curvature_estimator.BlockDiagonalCurvature
	ExplicitExactCurvature = curvature_estimator.ExplicitExactCurvature
	ImplicitExactCurvature = curvature_estimator.ImplicitExactCurvature
	set_default_tag_to_block_ctor = (
	curvature_estimator.set_default_tag_to_block_ctor)
	get_default_tag_to_block_ctor = (
	curvature_estimator.get_default_tag_to_block_ctor)
	OptaxPreconditioner = curvature_estimator.OptaxPreconditioner
	OptaxPreconditionState = curvature_estimator.OptaxPreconditionState

	# Optimizers
	Optimizer = optimizer.Optimizer

	HAIKU_BIASES = optimizer.HAIKU_BIASES
	HAIKU_BIASES_AND_NORMS = optimizer.HAIKU_BIASES_AND_NORMS

	__all__ = (
	# Modules
	"utils",
	"patches_second_moment",
	"layers_and_loss_tags",
	"loss_functions",
	"tag_graph_matcher",
	"tracer",
	"curvature_blocks",
	"curvature_estimator",
	"optimizer",
	# Patches second moments
	"patches_moments",
	"patches_moments_explicit",
	# Layer and loss tags
	"LossTag",
	"LayerTag",
	"register_generic",
	"register_dense",
	"register_conv2d",
	"register_scale_and_shift",
	# Tag graph matcher
	"auto_register_tags",
	# Tracer
	"ProcessedJaxpr",
	"loss_tags_vjp",
	"loss_tags_jvp",
	"loss_tags_hvp",
	"layer_tags_vjp",
	# Loss functions
	"LossFunction",
	"NegativeLogProbLoss",
	"DistributionNegativeLogProbLoss",
	"NormalMeanNegativeLogProbLoss",
	"NormalMeanVarianceNegativeLogProbLoss",
	"MultiBernoulliNegativeLogProbLoss",
	"CategoricalLogitsNegativeLogProbLoss",
	"OneHotCategoricalLogitsNegativeLogProbLoss",
	"register_sigmoid_cross_entropy_loss",
	"register_multi_bernoulli_predictive_distribution",
	"register_softmax_cross_entropy_loss",
	"register_categorical_predictive_distribution",
	"register_squared_error_loss",
	"register_normal_predictive_distribution",
	# Curvature blocks
	"CurvatureBlock",
	"ScaledIdentity",
	"Diagonal",
	"Full",
	"KroneckerFactored",
	"NaiveDiagonal",
	"NaiveFull",
	"NaiveTNT",
	"DenseDiagonal",
	"DenseFull",
	"DenseTwoKroneckerFactored",
	"RepeatedDenseKroneckerFactored",
	"DenseTNT",
	"Conv2DDiagonal",
	"Conv2DFull",
	"Conv2DTwoKroneckerFactored",
	"Conv2DTNT",
	"ScaleAndShiftDiagonal",
	"ScaleAndShiftFull",
	"set_max_parallel_elements",
	"get_max_parallel_elements",
	"set_default_eigen_decomposition_threshold",
	"get_default_eigen_decomposition_threshold",
	# Estimators
	"CurvatureEstimator",
	"BlockDiagonalCurvature",
	"ExplicitExactCurvature",
	"ImplicitExactCurvature",
	"set_default_tag_to_block_ctor",
	"get_default_tag_to_block_ctor",
	# Optimizers
	"Optimizer",
	)

	# _________________________________________
	# / Please don't use symbols in `_src` they \
	# \ are not part of the KFAC Jax public API./
	# -----------------------------------------
	# \ ^__^
	# \ (oo)\_______
	# (__)\ )\/\
	# \|\|----w \|
	# \|\| \|\|
	#
	try:
	del _src # pylint: disable=undefined-variable
	except NameError:
	pass