# Copyright 2022 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
#
#      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.
"""Base abstract class for all numerics."""

import abc
import flax.struct


class QNumerics(flax.struct.PyTreeNode, abc.ABC):
	"""Numerics for int8, int4, binary, etc."""

	# it in test. Remove and leave only get_dtype(

	@abc.abstractmethod
	def get_dtype(self):
		pass

	@abc.abstractmethod
	def fwd(self, x, context):
		"""Forward pass."""
		pass

	@abc.abstractmethod
	def abs_val_mapped_to(self):
		"""The value returned is the end of quantization range.

		It could be the biggest value that can be represented by numerical format
		exactly. E.g. in case of int8, 127 . Or it could be edge of the last bucket.
		Edge in case of int8, 127.5
		"""
		pass

	@abc.abstractmethod
	def vjp_fwd(self, x, context):
		pass

	@abc.abstractmethod
	def vjp_bwd(self, res, grad):
		pass
