layer is an operation with state takes care of weights variables and tracking thereof deferred build trainable argument value inference masking autocasting and some more magic addloss metric tracking rng seed tracking activity regularization this is the class from which all layers inherit a layer is a callable object that takes as input one or more tensors and that outputs one or more tensors it involves computation defined in the call method and a state weight variables state can be created in init for instance via self addweight in the optional build method which is invoked by the first call to the layer and supplies the shapes of the inputs which may not have been known at initialization time layers are recursively composable if you assign a layer instance as an attribute of another layer the outer layer will start tracking the weights the inner layer nested layers should be instantiated in the init method or build method users will just instantiate a layer and then treat it as a callable args trainable boolean whether the layer s variables should be trainable name string name of the layer dtype the dtype of the layer s computations and weights can also be a keras mixedprecision dtypepolicy which allows the computation and weight dtype to differ defaults to none none means to use keras mixedprecision dtypepolicy which is a float32 policy unless set to different value via keras mixedprecision setdtypepolicy attributes name the name of the layer string dtype dtype of the layer s weights alias of layer variabledtype variabledtype dtype of the layer s weights computedtype the dtype of the layer s computations layers automatically cast inputs to this dtype which causes the computations and output to also be in this dtype when mixed precision is used with a keras mixedprecision dtypepolicy this will be different than variabledtype trainableweights list of variables to be included in backprop nontrainableweights list of variables that should not be included in backprop weights the concatenation of the lists trainableweights and nontrainableweights in this order trainable whether the layer should be trained boolean i e whether its potentiallytrainable weights should be returned as part of layer trainableweights inputspec optional list of inputspec objects specifying the constraints on inputs that can be accepted by the layer we recommend that descendants of layer implement the following methods init defines custom layer attributes and creates layer weights that do not depend on input shapes using addweight or other state buildself inputshape this method can be used to create weights that depend on the shapes of the inputs using addweight or other state call will automatically build the layer if it has not been built yet by calling build callself args kwargs called in call after making sure build has been called call performs the logic of applying the layer to the input arguments two reserved keyword arguments you can optionally use in call are 1 training boolean whether the call is in inference mode or training mode 2 mask boolean tensor encoding masked timesteps in the input used e g in rnn layers a typical signature for this method is callself inputs and user could optionally add training and mask if the layer need them getconfigself returns a dictionary containing the configuration used to initialize this layer if the keys differ from the arguments in init then override fromconfigself as well this method is used when saving the layer or a model that contains this layer examples here s a basic example a layer with two variables w and b that returns y w x b it shows how to implement build and call variables set as attributes of a layer are tracked as weights of the layers in layer weights python class simpledenselayer def initself units32 super init self units units create the state of the layer weights def buildself inputshape self kernel self addweight shapeinputshape1 self units initializerglorotuniform trainabletrue namekernel self bias self addweight shapeself units initializerzeros trainabletrue namebias defines the computation def callself inputs return ops matmulinputs self kernel self bias instantiates the layer linearlayer simpledense4 this will also call buildinputshape and create the weights y linearlayerops ones2 2 assert lenlinearlayer weights 2 these weights are trainable so they re listed in trainableweights assert lenlinearlayer trainableweights 2 besides trainable weights updated via backpropagation during training layers can also have nontrainable weights these weights are meant to be updated manually during call here s a example layer that computes the running sum of its inputs python class computesumlayer def initself inputdim supercomputesum self init create a nontrainable weight self total self addweight shape initializerzeros trainablefalse nametotal def callself inputs self total assignself total ops suminputs return self total mysum computesum2 x ops ones2 2 y mysumx assert mysum weights mysum total assert mysum nontrainableweights mysum total assert mysum trainableweights wrap the userprovided build method in the builddecorator to add name scope support and serialization support record build config set built post build actions and lock state whether to automatically convert autocast inputs to call whether to allow nontensors as positional arguments in call dict of shapes that were used to call build remove attribute tracking for lists tfspecific attribute reset attribute tracking tfspecific prevent further state updates called automatically in build if not self tracker locked self tracker lock msg you cannot add new elements of state variables or sublayers to a layer that is already built all state must be created in the init method or in the build method def getbuildconfigself if self buildshapesdict is not none if lenself buildshapesdict 1 return inputshape tupleself buildshapesdict values0 else return shapesdict self buildshapesdict def buildfromconfigself config if config if inputshape in config self buildconfiginputshape elif shapesdict in config self buildconfigshapesdict self built true def addvariable self shape initializer dtypenone trainabletrue regularizernone constraintnone namenone return self addweight shapeshape initializerinitializer dtypedtype trainabletrainable regularizerregularizer constraintconstraint namename def addweight self shapenone initializernone dtypenone trainabletrue regularizernone constraintnone namenone self checksupercalled if shape is none shape if dtype is not none dtype backend standardizedtypedtype else dtype self variabledtype if initializer is none if float in dtype initializer glorotuniform else initializer zeros initializer initializers getinitializer with backend namescopeself name callerself variable backend variable initializerinitializer shapeshape dtypedtype trainabletrainable namename will be added to layer losses variable regularizer regularizer variable constraint constraint self trackvariablevariable return variable property def trainableself sets trainable attribute for the layer and its sublayers when this value is changed during training e g with a callback you need to call the parent model maketrainfunction with forcetrue in order to recompile the training graph args value boolean with the desired state for the layer s trainable attribute list of all layer state including random seeds this extends layer weights to include all state used by the layer including seedgenerators note that metrics variables are not included here use metricsvariables to visit all the metric variables return all variables associate with the layer including metrics and random seeds also deduplicate them list of all trainable layer state this is equivalent to layer trainableweights list of all nontrainable layer state this extends layer nontrainableweights to include all state used by the layer including state for metrics and seedgenerators list of all weight variables of the layer unlike layer variables this excludes metric state and random seeds return only variables directly owned by layers and sublayers also deduplicate them list of all trainable weight variables of the layer these are the weights that get updated by the optimizer during training list of all nontrainable weight variables of the layer these are the weights that should not be updated by the optimizer during training unlike layer nontrainablevariables this excludes metric state and random seeds list of all metric variables vars for metric in self metrics vars extendmetric variables for layer in self layers for metric in layer metrics vars extendmetric variables return vars def getweightsself sets the values of layer weights from a list of numpy arrays layerweights self weights if lenlayerweights lenweights raise valueerror fyou called setweightsweights on layer self name fwith a weight list of length lenweights but the layer fwas expecting lenlayerweights weights for variable value in ziplayerweights weights if variable shape value shape raise valueerror flayer self name weight shape variable shape is not compatible with provided weight fshape value shape variable assignvalue property def dtypeself the dtype of the computations performed by the layer return self dtypepolicy computedtype property def variabledtypeself the dtype layer inputs should be converted to return self dtypepolicy computedtype property def supportsmaskingself 1 convert any array arguments to tensors of correct dtype used to avoid expensive tree operations in the most common case 2 enforce that only tensors can be passed positionally caches info about call signature args kwargs 3 check input spec for 1st positional arg todo consider extending this to all args and kwargs 4 call build 5 infer training value training phase for layer call is set via in order of priority 1 the training argument passed to this layer call if not none 2 the training argument of an outer layer call 4 any nonnone default value for training in the call signature 5 false treating the layer as if it s in inference maintains info about the layer call stack across nested calls this is the value explicity passed by the user wasn t passed explicitly use context value get signature default value only populate arg if it has a concrete value 6 populate mask arguments 7 call the layer clear or update the current scope if necessary some preprocessing layers might have a nonfloat dtype we should not autocast in this case enter a new scope if our dtypes are mixed change the layout for the layer output if needed this is useful for relayout intermediate tensor in the model to achieve the optimal performance record activity regularizer loss set masks on outputs provided only the first positional input arg and its mask todo consider extending this to all args and kwargs destroy call context if we created it call the layer without any side effects args trainablevariables list of trainable variables of the model nontrainablevariables list of nontrainable variables of the model args positional argumets to be passed to call returnlosses if true statelesscall will return the list of losses created during call as part of its return values kwargs keyword arguments to be passed to call returns a tuple by default returns outputs nontrainablevariables if returnlosses true then returns outputs nontrainablevariables losses note nontrainablevariables include not only nontrainable weights such as batchnormalization statistics but also rng seed state if there are any random operations part of the layer such as dropout and metric state if there are any metrics attached to the layer these are all elements of state of the layer example python model data trainablevariables model trainablevariables nontrainablevariables model nontrainablevariables call the model with zero side effects outputs nontrainablevariables model statelesscall trainablevariables nontrainablevariables data attach the updated state to the model until you do this the model is still in its precall state for refvar value in zip model nontrainablevariables nontrainablevariables refvar assignvalue gather variable mapping call in stateless scope gather updated nontrainable variables use computeoutputshape to return the right output spec case nested could be a tuplelist of shapes or a dict of shapes could be deeply nested can be called inside of the call method to add a scalar loss example python class mylayerlayer def callself x self addlossops sumx return x eager only list of scalar losses from addloss regularizers and sublayers losses self getownlosses for layer in self flattenlayersincludeselffalse losses extendlayer getownlosses weightregularizationlosses for v in self trainableweights if backend instatelessscope v backend getstatelessscope getcurrentvaluev regularizer getattrv regularizer none if regularizer weightregularizationlosses appendregularizerv losses extendweightregularizationlosses return losses def clearlossesself if backend instatelessscope scope backend getstatelessscope if scope collectlosses for x in scope losses if idx in self lossids scope losses removex self losses clear self lossids clear for layer in self layers layer clearlosses def saveownvariablesself store allvars self trainablevariables self nontrainablevariables for i v in enumerateallvars storefi v numpy def loadownvariablesself store allvars self trainablevariables self nontrainablevariables if lenstore keys lenallvars if lenallvars 0 and not self built raise valueerror flayer self name was never built and thus it doesn t have any variables fhowever the weights file lists lenstore keys variables for this layer n in most cases this error indicates that either nn 1 the layer is owned by a parent layer that implements a build method but calling the parent s build method did not create the state of fthe child layer self name a build method must create all state for the layer including the state of any children layers nn 2 you need to implement the def buildfromconfigself config method fon layer self name to specify how to rebuild it during loading in this case you might also want to implement the method that generates the build config at saving time def getbuildconfigself the method buildfromconfig is meant to create the state of the layer i e its variables upon deserialization raise valueerror flayer self name expected lenallvars variables but received flenstore keys variables during loading fexpected v name for v in allvars for i v in enumerateallvars v assignstorefi def trackvariableself variable if variable trainable self tracker addtostoretrainablevariables variable else self tracker addtostorenontrainablevariables variable def addmetricself permanently disabled raise notimplementederror def countparamsself if not self built raise valueerror you tried to call countparams fon layer self name but the layer isn t built you can build it manually via flayer buildinputshape return summaryutils countparamsself weights def maybebuildself callspec if self built return shapesdict getshapesdictcallspec firstshape nextitershapesdict values none if the layer has a build method call it with our input shapes if not utils isdefaultself build shapesdict updateshapesdictfortargetfn self build shapesdictshapesdict callspeccallspec classnameself class name self buildshapesdict check input spec again after build since self inputspec may have been updated self assertinputcompatibilitycallspec firstarg return otherwise attempt to build the layer by calling it on symbolic input if mighthaveunbuiltstateself try backend computeoutputspec self call callspec argumentsdict except exception as e if callspec eager will let the actual eager call do statebuilding return warnings warn flayer self name looks like it has unbuilt state but keras is not able to trace the layer call in order to build it automatically possible causes n 1 the call method of your layer may be crashing try to call the layer eagerly on some test input first to see if it works e g x np random random3 4 y layerxn 2 if the call method is correct then you may need to implement the def buildself inputshape method on your layer it should create all variables used by the layer e g by calling layer build on all its children layers n fexception encoutered e self buildfirstshape def buildbyrunforsingleposargself inputshape case all inputs are in the first arg possibly nested inputtensors mapshapestructure lambda s backend kerastensors inputshape try backend computeoutputspecself call inputtensors return true except return false def buildbyrunforkwargsself shapesdict case inputs were recorded as multiple keyword arguments if allisshapetuples for s in shapesdict values case all input keyword arguments were plain tensors inputtensors we strip the shape suffix to recover kwarg names utils removesuffixk shape backend kerastensorshape for k shape in shapesdict items try backend computeoutputspecself call inputtensors return true except return false else not supported nested input keyword arguments return false def reprself return fself class name fnameself name builtself built def strself return fself class name fnameself name builtself built def setattrself name value track variables layers metrics seedgenerators name value self setattrhookname value if hasattrself tracker value self tracker trackvalue elif name tracker self initializertracker return super setattrname value def checksupercalledself if getattrself lock true raise runtimeerror fin layer self class name you forgot to call super init as the first statement in the init method go add it def assertinputcompatibilityself arg0 if self inputspec inputspec assertinputcompatibility self inputspec arg0 layernameself name def getcallcontextself enter new call context introspect recursively through sublayers numpy backend does not support masking it s a c type training and mask are special kwargs that are always available in a layer if user specifies them in their call without adding to spec we remove them to be able to bind variables user is not using training anyway so we can ignore todo if necessary use workaround for mask return a dict mapping argument names to their values sig inspect signaturefn boundargs sig bindargs kwargs argdict for name value in boundargs arguments items argdictname value return argdict def getshapesdictcallspec shapesdict for k v in callspec tensorargumentsdict items if k mask or k endswithmask do not include mask tensors in shapes dict continue if k kwargs or k args do not include catchalls in shapes dict continue if k in callspec nestedtensorargumentnames shapesdictfkshape tree mapstructure lambda x backend standardizeshapex shape v else shapesdictfkshape backend standardizeshapev shape return shapesdict def updateshapesdictfortargetfn targetfn shapesdict callspec classname if utils isdefaulttargetfn return none sig inspect signaturetargetfn expectednames for name param in sig parameters items if param kind in param positionalorkeyword param positionalonly param keywordonly expectednames appendname single arg don t check names pass first shape if lenexpectednames 1 key expectednames0 values tupleshapesdict values if values inputshape values0 else inputshape none return key inputshape multiple args check that all names line up kwargs for name in expectednames methodname targetfn name errorpreamble ffor a methodname method with more than one argument all arguments should have a shape suffix and match an argument ffrom call e g methodnameself fooshape barshape if not name endswithshape raise valueerror ferrorpreamble for layer classname freceived methodname argument fname which does not end in shape expectedcallarg utils removesuffixname shape if expectedcallarg not in callspec argumentsdict raise valueerror ferrorpreamble for layer classname freceived methodname argument fname but call does not have argument fexpectedcallarg if name in shapesdict kwargsname shapesdictname return kwargs class callcontext def initself entrylayer self entrylayer entrylayer self training none def isshapetuples return isinstances list tuple and all d is none or isinstanced int for d in s def mighthaveunbuiltstatelayer return anynot lr built for lr in layer layers this is the class from which all layers inherit a layer is a callable object that takes as input one or more tensors and that outputs one or more tensors it involves computation defined in the call method and a state weight variables state can be created in __init__ for instance via self add_weight in the optional build method which is invoked by the first __call__ to the layer and supplies the shape s of the input s which may not have been known at initialization time layers are recursively composable if you assign a layer instance as an attribute of another layer the outer layer will start tracking the weights the inner layer nested layers should be instantiated in the __init__ method or build method users will just instantiate a layer and then treat it as a callable args trainable boolean whether the layer s variables should be trainable name string name of the layer dtype the dtype of the layer s computations and weights can also be a keras mixed_precision dtypepolicy which allows the computation and weight dtype to differ defaults to none none means to use keras mixed_precision dtype_policy which is a float32 policy unless set to different value via keras mixed_precision set_dtype_policy attributes name the name of the layer string dtype dtype of the layer s weights alias of layer variable_dtype variable_dtype dtype of the layer s weights compute_dtype the dtype of the layer s computations layers automatically cast inputs to this dtype which causes the computations and output to also be in this dtype when mixed precision is used with a keras mixed_precision dtypepolicy this will be different than variable_dtype trainable_weights list of variables to be included in backprop non_trainable_weights list of variables that should not be included in backprop weights the concatenation of the lists trainable_weights and non_trainable_weights in this order trainable whether the layer should be trained boolean i e whether its potentially trainable weights should be returned as part of layer trainable_weights input_spec optional list of inputspec object s specifying the constraints on inputs that can be accepted by the layer we recommend that descendants of layer implement the following methods __init__ defines custom layer attributes and creates layer weights that do not depend on input shapes using add_weight or other state build self input_shape this method can be used to create weights that depend on the shape s of the input s using add_weight or other state __call__ will automatically build the layer if it has not been built yet by calling build call self args kwargs called in __call__ after making sure build has been called call performs the logic of applying the layer to the input arguments two reserved keyword arguments you can optionally use in call are 1 training boolean whether the call is in inference mode or training mode 2 mask boolean tensor encoding masked timesteps in the input used e g in rnn layers a typical signature for this method is call self inputs and user could optionally add training and mask if the layer need them get_config self returns a dictionary containing the configuration used to initialize this layer if the keys differ from the arguments in __init__ then override from_config self as well this method is used when saving the layer or a model that contains this layer examples here s a basic example a layer with two variables w and b that returns y w x b it shows how to implement build and call variables set as attributes of a layer are tracked as weights of the layers in layer weights python class simpledense layer def __init__ self units 32 super __init__ self units units create the state of the layer weights def build self input_shape self kernel self add_weight shape input_shape 1 self units initializer glorot_uniform trainable true name kernel self bias self add_weight shape self units initializer zeros trainable true name bias defines the computation def call self inputs return ops matmul inputs self kernel self bias instantiates the layer linear_layer simpledense 4 this will also call build input_shape and create the weights y linear_layer ops ones 2 2 assert len linear_layer weights 2 these weights are trainable so they re listed in trainable_weights assert len linear_layer trainable_weights 2 besides trainable weights updated via backpropagation during training layers can also have non trainable weights these weights are meant to be updated manually during call here s a example layer that computes the running sum of its inputs python class computesum layer def __init__ self input_dim super computesum self __init__ create a non trainable weight self total self add_weight shape initializer zeros trainable false name total def call self inputs self total assign self total ops sum inputs return self total my_sum computesum 2 x ops ones 2 2 y my_sum x assert my_sum weights my_sum total assert my_sum non_trainable_weights my_sum total assert my_sum trainable_weights wrap the user provided build method in the build_decorator to add name scope support and serialization support record build config set built post build actions and lock state whether to automatically convert auto cast inputs to call whether to allow non tensors as positional arguments in call dict of shapes that were used to call build remove attribute tracking for lists tf specific attribute reset attribute tracking tf specific prevent further state updates called automatically in build returns a dictionary with the layer s input shape this method returns a config dict that can be used by build_from_config config to create all states e g variables and lookup tables needed by the layer by default the config only contains the input shape that the layer was built with if you re writing a custom layer that creates state in an unusual way you should override this method to make sure this state is already created when keras attempts to load its value upon model loading returns a dict containing the input shape associated with the layer builds the layer s states with the supplied config dict by default this method calls the build config input_shape method which creates weights based on the layer s input shape in the supplied config if your config contains other information needed to load the layer s state you should override this method args config dict containing the input shape associated with this layer add a weight variable to the layer alias of add_weight add a weight variable to the layer args shape shape tuple for the variable must be fully defined no none entries defaults to scalar if unspecified initializer initializer object to use to populate the initial variable value or string name of a built in initializer e g random_normal if unspecified defaults to glorot_uniform for floating point variables and to zeros for all other types e g int bool dtype dtype of the variable to create e g float32 if unspecified defaults to the layer s variable dtype which itself defaults to float32 if unspecified trainable boolean whether the variable should be trainable via backprop or whether its updates are managed manually constraint contrainst object to call on the variable after any optimizer update or string name of a built in constraint name string name of the variable useful for debugging purposes will be added to layer losses settable boolean whether this layer should be trainable or not sets trainable attribute for the layer and its sublayers when this value is changed during training e g with a callback you need to call the parent model make_train_function with force true in order to recompile the training graph args value boolean with the desired state for the layer s trainable attribute list of all layer state including random seeds this extends layer weights to include all state used by the layer including seedgenerator s note that metrics variables are not included here use metrics_variables to visit all the metric variables return all variables associate with the layer including metrics and random seeds also deduplicate them list of all trainable layer state this is equivalent to layer trainable_weights list of all non trainable layer state this extends layer non_trainable_weights to include all state used by the layer including state for metrics and seedgenerator s list of all weight variables of the layer unlike layer variables this excludes metric state and random seeds return only variables directly owned by layers and sub layers also deduplicate them list of all trainable weight variables of the layer these are the weights that get updated by the optimizer during training list of all non trainable weight variables of the layer these are the weights that should not be updated by the optimizer during training unlike layer non_trainable_variables this excludes metric state and random seeds list of all metric variables return the values of layer weights as a list of numpy arrays sets the values of layer weights from a list of numpy arrays alias of layer variable_dtype the dtype of the computations performed by the layer the dtype of the state weights of the layer the dtype layer inputs should be converted to whether this layer supports computing a mask using compute_mask 1 convert any array arguments to tensors of correct dtype used to avoid expensive tree operations in the most common case 2 enforce that only tensors can be passed positionally caches info about call signature args kwargs 3 check input spec for 1st positional arg todo consider extending this to all args and kwargs 4 call build 5 infer training value training phase for layer call is set via in order of priority 1 the training argument passed to this layer call if not none 2 the training argument of an outer layer call 4 any non none default value for training in the call signature 5 false treating the layer as if it s in inference maintains info about the layer call stack across nested calls this is the value explicity passed by the user wasn t passed explicitly use context value get signature default value only populate arg if it has a concrete value 6 populate mask argument s 7 call the layer clear or update the current scope if necessary some preprocessing layers might have a non float dtype we should not autocast in this case enter a new scope if our dtypes are mixed change the layout for the layer output if needed this is useful for relayout intermediate tensor in the model to achieve the optimal performance record activity regularizer loss set masks on outputs provided only the first positional input arg and its mask todo consider extending this to all args and kwargs destroy call context if we created it call the layer without any side effects args trainable_variables list of trainable variables of the model non_trainable_variables list of non trainable variables of the model args positional argumets to be passed to call return_losses if true stateless_call will return the list of losses created during call as part of its return values kwargs keyword arguments to be passed to call returns a tuple by default returns outputs non_trainable_variables if return_losses true then returns outputs non_trainable_variables losses note non_trainable_variables include not only non trainable weights such as batchnormalization statistics but also rng seed state if there are any random operations part of the layer such as dropout and metric state if there are any metrics attached to the layer these are all elements of state of the layer example python model data trainable_variables model trainable_variables non_trainable_variables model non_trainable_variables call the model with zero side effects outputs non_trainable_variables model stateless_call trainable_variables non_trainable_variables data attach the updated state to the model until you do this the model is still in its pre call state for ref_var value in zip model non_trainable_variables non_trainable_variables ref_var assign value gather variable mapping call in stateless scope gather updated non trainable variables use compute_output_shape to return the right output spec case nested could be a tuple list of shapes or a dict of shapes could be deeply nested can be called inside of the call method to add a scalar loss example python class mylayer layer def call self x self add_loss ops sum x return x eager only list of scalar losses from add_loss regularizers and sublayers saves the state of the layer you can override this method to take full control of how the state of the layer is saved upon calling model save args store dict where the state of the model will be saved loads the state of the layer you can override this method to take full control of how the state of the layer is loaded upon calling keras models load_model args store dict from which the state of the model will be loaded permanently disabled count the total number of scalars composing the weights returns an integer count if the layer has a build method call it with our input shapes check input spec again after build since self input_spec may have been updated otherwise attempt to build the layer by calling it on symbolic input will let the actual eager call do state building case all inputs are in the first arg possibly nested case inputs were recorded as multiple keyword arguments case all input keyword arguments were plain tensors we strip the _shape suffix to recover kwarg names not supported nested input keyword arguments track variables layers metrics seedgenerators returns currently active callcontext enter new call context introspect recursively through sublayers numpy backend does not support masking it s a c type training and mask are special kwargs that are always available in a layer if user specifies them in their call without adding to spec we remove them to be able to bind variables user is not using training anyway so we can ignore todo if necessary use workaround for mask return a dict mapping argument names to their values convert the call arguments dict into a dict of input shape arguments example get_shapes_dict call_spec input_a_shape 2 3 do not include mask tensors in shapes dict do not include catch alls in shapes dict updates a shapes_dict for build or compute_output_shape this function will align a dictionary of the shapes of all tensor passed to call with the signatures of build or compute_output_shape the alignment is a follows if build or compute_output_shape accept only one argument forward the shape of the first positional argument from call without checking any argument names if build or compute_output_shape accept multiple arguments enforce that all argument names match a call argument name e g foo_shape would match call argument foo returns an updated shapes_dict that can be used to invoke target_fn shapes_dict single arg don t check names pass first shape multiple args check that all names line up
import collections
import inspect
import warnings
from functools import wraps

import tree

from keras import backend
from keras import initializers
from keras import mixed_precision
from keras import regularizers
from keras import utils
from keras.api_export import keras_export
from keras.backend import KerasTensor
from keras.backend.common import global_state
from keras.backend.common.name_scope import current_path
from keras.distribution import distribution_lib
from keras.layers import input_spec
from keras.metrics.metric import Metric
from keras.ops.operation import Operation
from keras.utils import python_utils
from keras.utils import summary_utils
from keras.utils import traceback_utils
from keras.utils import tracking
from keras.utils.shape_utils import map_shape_structure

if backend.backend() == "tensorflow":
    from keras.backend.tensorflow.layer import TFLayer as BackendLayer
elif backend.backend() == "jax":
    from keras.backend.jax.layer import JaxLayer as BackendLayer
elif backend.backend() == "torch":
    from keras.backend.torch.layer import TorchLayer as BackendLayer
elif backend.backend() == "numpy":
    from keras.backend.numpy.layer import NumpyLayer as BackendLayer
else:
    raise RuntimeError(
        f"Backend '{backend.backend()}' must implement a layer mixin class."
    )


@keras_export(["keras.Layer", "keras.layers.Layer"])
class Layer(BackendLayer, Operation):
    

    def __new__(cls, *args, **kwargs):
        
        
        obj = super().__new__(cls, *args, **kwargs)

        original_build_method = obj.build

        @wraps(original_build_method)
        def build_wrapper(*args, **kwargs):
            with backend.name_scope(obj.name, caller=obj):
                original_build_method(*args, **kwargs)
            
            signature = inspect.signature(original_build_method)
            obj._build_shapes_dict = signature.bind(*args, **kwargs).arguments
            
            obj.built = True
            obj._post_build()
            obj._lock_state()

        obj.build = build_wrapper
        return obj

    def __init__(
        self,
        *,
        activity_regularizer=None,
        trainable=True,
        dtype=None,
        autocast=True,
        name=None,
        **kwargs,
    ):
        BackendLayer.__init__(self)
        self._lock = False
        Operation.__init__(self, name=name)
        self.activity_regularizer = regularizers.get(activity_regularizer)
        input_dim_arg = kwargs.pop("input_dim", None)
        if input_dim_arg is not None:
            input_shape_arg = (input_dim_arg,)
        else:
            input_shape_arg = kwargs.pop("input_shape", None)
        if input_shape_arg is not None:
            warnings.warn(
                "Do not pass an `input_shape`/`input_dim` argument to "
                "a layer. When using Sequential models, "
                "prefer using an `Input(shape)` object as the "
                "first layer in the model instead.",
                stacklevel=2,
            )
            self._input_shape_arg = input_shape_arg
        if kwargs:
            raise ValueError(
                "Unrecognized keyword arguments "
                f"passed to {self.__class__.__name__}: {kwargs}"
            )

        self.built = False
        self.dtype_policy = mixed_precision.resolve_policy(dtype)
        self.autocast = autocast
        self._input_spec = None
        self._called = False
        self.supports_jit = True

        self._trainable = trainable
        self._losses = []
        self._loss_ids = set()

        self._call_signature = inspect.signature(self.call)
        call_signature_parameters = [
            p.name for p in self._call_signature.parameters.values()
        ]
        self._call_has_training_arg = "training" in call_signature_parameters
        self._call_has_mask_arg = "mask" in call_signature_parameters

        self._supports_masking = not utils.is_default(self.compute_mask)
        
        self._convert_input_args = True
        
        self._allow_non_tensor_positional_args = False
        
        self._build_shapes_dict = None
        self._initializer_tracker()

    @tracking.no_automatic_dependency_tracking
    def _initializer_tracker(self):
        if hasattr(self, "_tracker"):
            return

        trainable_variables = []
        non_trainable_variables = []
        layers = []
        metrics = []
        seed_generators = []
        self._tracker = tracking.Tracker(
            {
                "trainable_variables": (
                    lambda x: isinstance(x, backend.Variable) and x.trainable,
                    trainable_variables,
                ),
                "non_trainable_variables": (
                    lambda x: isinstance(x, backend.Variable)
                    and not x.trainable,
                    non_trainable_variables,
                ),
                "metrics": (lambda x: isinstance(x, Metric), metrics),
                "layers": (
                    lambda x: isinstance(x, Layer)
                    and not isinstance(x, Metric),
                    layers,
                ),
                "seed_generators": (
                    lambda x: isinstance(x, backend.random.SeedGenerator),
                    seed_generators,
                ),
            }
        )
        if backend.backend() == "tensorflow":
            
            _self_setattr_tracking = getattr(
                self, "_self_setattr_tracking", True
            )
            self._self_setattr_tracking = False

        self._trainable_variables = trainable_variables
        self._non_trainable_variables = non_trainable_variables
        self._layers = layers
        self._metrics = metrics
        self._seed_generators = seed_generators

        if backend.backend() == "tensorflow":
            
            self._self_setattr_tracking = _self_setattr_tracking

    @property
    def input_spec(self):
        return self._input_spec

    @input_spec.setter
    def input_spec(self, value):
        self._input_spec = value

    @utils.default
    def build(self, input_shape):
        self._check_super_called()
        if utils.is_default(self.build) and might_have_unbuilt_state(self):
            warnings.warn(
                f"`build()` was called on layer '{self.name}', however "
                "the layer does not have a `build()` method implemented "
                "and it looks like it has unbuilt state. This will cause "
                "the layer to be marked as built, despite not being "
                "actually built, which may cause failures down the line. "
                "Make sure to implement a proper `build()` method."
            )
        self.built = True

    def _lock_state(self):
        
        if not self._tracker.locked:
            self._tracker.lock(
                msg=(
                    "You cannot add new elements of state "
                    "(variables or sub-layers) "
                    "to a layer that is already built. All state "
                    "must be created in the `__init__()` method or "
                    "in the `build()` method."
                )
            )

    def get_build_config(self):
        
        if self._build_shapes_dict is not None:
            if len(self._build_shapes_dict) == 1:
                return {
                    "input_shape": tuple(self._build_shapes_dict.values())[0],
                }
            else:
                return {"shapes_dict": self._build_shapes_dict}

    def build_from_config(self, config):
        
        if config:
            if "input_shape" in config:
                self.build(config["input_shape"])
            elif "shapes_dict" in config:
                self.build(**config["shapes_dict"])
            self.built = True

    def add_variable(
        self,
        shape,
        initializer,
        dtype=None,
        trainable=True,
        regularizer=None,
        constraint=None,
        name=None,
    ):
        
        return self.add_weight(
            shape=shape,
            initializer=initializer,
            dtype=dtype,
            trainable=trainable,
            regularizer=regularizer,
            constraint=constraint,
            name=name,
        )

    def add_weight(
        self,
        shape=None,
        initializer=None,
        dtype=None,
        trainable=True,
        regularizer=None,
        constraint=None,
        name=None,
    ):
        
        self._check_super_called()
        if shape is None:
            shape = ()
        if dtype is not None:
            dtype = backend.standardize_dtype(dtype)
        else:
            dtype = self.variable_dtype
        if initializer is None:
            if "float" in dtype:
                initializer = "glorot_uniform"
            else:
                initializer = "zeros"
        initializer = initializers.get(initializer)
        with backend.name_scope(self.name, caller=self):
            variable = backend.Variable(
                initializer=initializer,
                shape=shape,
                dtype=dtype,
                trainable=trainable,
                name=name,
            )
        
        variable.regularizer = regularizer
        variable.constraint = constraint
        self._track_variable(variable)
        return variable

    @property
    def trainable(self):
        
        return self._trainable

    @trainable.setter
    def trainable(self, value):
        
        value = bool(value)
        self._trainable = value
        for v in self._trainable_variables:
            v.trainable = value
        for layer in self._layers:
            layer.trainable = value

    @property
    def variables(self):
        
        
        
        variables = []
        seen_ids = set()
        for v in self._trainable_variables + self._non_trainable_variables:
            if id(v) not in seen_ids:
                variables.append(v)
                seen_ids.add(id(v))
        for sg in self._seed_generators:
            variables.append(sg.state)
        for layer in self._layers:
            for v in layer.variables:
                if id(v) not in seen_ids:
                    variables.append(v)
                    seen_ids.add(id(v))
        return variables

    @property
    def trainable_variables(self):
        
        if not self.trainable:
            return []
        return [v for v in self.variables if v.trainable]

    @property
    def non_trainable_variables(self):
        
        if not self.trainable:
            return self.variables
        return [v for v in self.variables if not v.trainable]

    @property
    def weights(self):
        
        
        
        weights = []
        seen_ids = set()
        for w in self._trainable_variables + self._non_trainable_variables:
            if id(w) not in seen_ids:
                weights.append(w)
                seen_ids.add(id(w))
        for layer in self._layers:
            for w in layer.weights:
                if id(w) not in seen_ids:
                    weights.append(w)
                    seen_ids.add(id(w))
        return weights

    @property
    def trainable_weights(self):
        
        if not self.trainable:
            return []
        return [v for v in self.weights if v.trainable]

    @property
    def non_trainable_weights(self):
        
        if not self.trainable:
            return self.weights
        return [v for v in self.weights if not v.trainable]

    @property
    def metrics_variables(self):
        
        vars = []
        for metric in self._metrics:
            vars.extend(metric.variables)
        for layer in self._layers:
            for metric in layer._metrics:
                vars.extend(metric.variables)
        return vars

    def get_weights(self):
        
        return [v.numpy() for v in self.weights]

    def set_weights(self, weights):
        
        layer_weights = self.weights
        if len(layer_weights) != len(weights):
            raise ValueError(
                f"You called `set_weights(weights)` on layer '{self.name}' "
                f"with a weight list of length {len(weights)}, but the layer "
                f"was expecting {len(layer_weights)} weights."
            )
        for variable, value in zip(layer_weights, weights):
            if variable.shape != value.shape:
                raise ValueError(
                    f"Layer {self.name} weight shape {variable.shape} "
                    "is not compatible with provided weight "
                    f"shape {value.shape}."
                )
            variable.assign(value)

    @property
    def dtype(self):
        
        return self.variable_dtype

    @property
    def compute_dtype(self):
        
        return self.dtype_policy.compute_dtype

    @property
    def variable_dtype(self):
        
        return self.dtype_policy.variable_dtype

    @property
    def input_dtype(self):
        
        return self.dtype_policy.compute_dtype

    @property
    def supports_masking(self):
        
        return self._supports_masking

    @supports_masking.setter
    def supports_masking(self, value):
        self._supports_masking = value

    @utils.default
    def compute_mask(self, inputs, previous_mask):
        return previous_mask

    @traceback_utils.filter_traceback
    def __call__(self, *args, **kwargs):
        self._check_super_called()
        self._called = True

        
        
        def maybe_convert(x):
            if backend.is_tensor(x):
                if (
                    self.autocast
                    and backend.is_float_dtype(x.dtype)
                    and x.dtype != self.input_dtype
                ):
                    x = backend.cast(x, dtype=self.input_dtype)
                return x
            elif isinstance(x, backend.KerasTensor):
                if (
                    self.autocast
                    and backend.is_float_dtype(x.dtype)
                    and x.dtype != self.input_dtype
                ):
                    x.dtype = self.input_dtype
                return x
            elif hasattr(x, "__array__"):
                return backend.convert_to_tensor(x, dtype=self.input_dtype)
            return x

        
        if (
            kwargs
            or len(args) != 1
            or not backend.is_tensor(args[0])
            or backend.standardize_dtype(args[0].dtype) != self.input_dtype
        ) and self._convert_input_args:
            args = tree.map_structure(maybe_convert, args)
            kwargs = tree.map_structure(maybe_convert, kwargs)

        
        
        if not self._allow_non_tensor_positional_args:
            for arg in tree.flatten(args):
                if not isinstance(arg, KerasTensor) and not backend.is_tensor(
                    arg
                ):
                    raise ValueError(
                        "Only input tensors may be passed as "
                        "positional arguments. The following argument value "
                        f"should be passed as a keyword argument: {arg} "
                        f"(of type {type(arg)})"
                    )

        
        call_spec = CallSpec(self._call_signature, args, kwargs)

        
        
        
        self._assert_input_compatibility(call_spec.first_arg)

        
        
        with backend.name_scope(self.name, caller=self):
            self._maybe_build(call_spec)

        
        
        
        
        
        
        

        
        
        call_context = self._get_call_context()

        
        training = call_spec.user_arguments_dict.get("training", None)
        if training is None:
            
            training = call_context.training
            if training is None:
                
                training = call_spec.arguments_dict.get("training", None)
        call_context.training = training
        if self._call_has_training_arg and training is not None:
            
            kwargs["training"] = training

        
        
        if len(call_spec.tensor_arguments_dict) == 1:
            if (
                "mask" in call_spec.argument_names
                and call_spec.arguments_dict["mask"] is None
            ):
                arg_name = list(call_spec.tensor_arguments_dict.keys())[0]
                only_tensor_arg = call_spec.tensor_arguments_dict[arg_name]
                mask = tree.map_structure(
                    lambda x: getattr(x, "_keras_mask", None),
                    only_tensor_arg,
                )
                kwargs["mask"] = mask
        elif len(call_spec.tensor_arguments_dict) > 1:
            for k, v in call_spec.tensor_arguments_dict.items():
                expected_mask_arg_name = f"{k}_mask"
                if expected_mask_arg_name in call_spec.argument_names:
                    if call_spec.arguments_dict[expected_mask_arg_name] is None:
                        mask = tree.map_structure(
                            lambda x: getattr(x, "_keras_mask", None), v
                        )
                        kwargs[expected_mask_arg_name] = mask

        
        
        try:
            with backend.name_scope(self.name, caller=self):
                current_scope = backend.get_autocast_scope()
                new_scope = None
                if current_scope is not None:
                    
                    if not self.autocast:
                        new_scope = backend.AutocastScope(None)
                    elif not backend.is_float_dtype(self.compute_dtype):
                        
                        
                        new_scope = backend.AutocastScope(None)
                    elif current_scope.dtype != self.compute_dtype:
                        new_scope = backend.AutocastScope(self.compute_dtype)
                elif self.compute_dtype != self.variable_dtype:
                    
                    new_scope = backend.AutocastScope(self.compute_dtype)

                if new_scope is not None:
                    with new_scope:
                        outputs = super().__call__(*args, **kwargs)
                else:
                    outputs = super().__call__(*args, **kwargs)
                
                
                
                distribution = distribution_lib.distribution()
                if distribution is not None:
                    current_layer_path = current_path()
                    current_layer_path += "/output"
                    layout = distribution.get_tensor_layout(current_layer_path)
                    if layout:
                        outputs = distribution_lib.distribute_tensor(
                            outputs, layout
                        )

                if not self.built:
                    self.built = True
                
                if self.activity_regularizer is not None:
                    for output in tree.flatten(outputs):
                        if backend.is_tensor(output):
                            self.add_loss(self.activity_regularizer(output))

            
            
            
            previous_mask = getattr(call_spec.first_arg, "_keras_mask", None)
            if self.supports_masking:
                self._set_mask_metadata(
                    call_spec.first_arg, outputs, previous_mask
                )
            elif previous_mask is not None:
                warnings.warn(
                    f"Layer '{self.name}' (of type {self.__class__.__name__}) "
                    "was passed an input with a mask attached to it. "
                    "However, this layer does not support masking and will "
                    "therefore destroy the mask information. Downstream "
                    "layers will not see the mask."
                )
        finally:
            
            self._maybe_reset_call_context()
        return outputs

    def call(self, *args, **kwargs):
        raise NotImplementedError(
            f"Layer {self.__class__.__name__} does not have a `call()` "
            "method implemented."
        )

    @traceback_utils.filter_traceback
    def stateless_call(
        self,
        trainable_variables,
        non_trainable_variables,
        *args,
        return_losses=False,
        **kwargs,
    ):
        
        self._check_super_called()

        if not self.built:
            raise ValueError(
                f"To call stateless_call, {self.__class__.__name__} must be "
                "built (i.e. its variables must have been already created). "
                "You can build it by calling it on some data."
            )
        if len(trainable_variables) != len(self.trainable_variables):
            raise ValueError(
                "Argument `trainable_variables` must be a list of tensors "
                "corresponding 1:1 to "
                f"{self.__class__.__name__}().trainable_variables. "
                f"Received list with length {len(trainable_variables)}, "
                f"but expected {len(self.trainable_variables)} variables."
            )
        if len(non_trainable_variables) != len(self.non_trainable_variables):
            raise ValueError(
                "Argument `non_trainable_variables` must be a list of tensors "
                "corresponding 1:1 to "
                f"{self.__class__.__name__}().non_trainable_variables. "
                f"Received list with length {len(non_trainable_variables)}, "
                f"but expected {len(self.non_trainable_variables)} variables."
            )

        
        trainable_mapping = zip(self.trainable_variables, trainable_variables)
        non_trainable_mapping = zip(
            self.non_trainable_variables, non_trainable_variables
        )
        mapping = list(trainable_mapping) + list(non_trainable_mapping)

        
        with backend.StatelessScope(
            state_mapping=mapping, collect_losses=return_losses
        ) as scope:
            outputs = self.call(*args, **kwargs)

        
        non_trainable_variables = []
        for v in self.non_trainable_variables:
            new_v = scope.get_current_value(v)
            if new_v is not None:
                non_trainable_variables.append(new_v)
            else:
                non_trainable_variables.append(v)

        if return_losses:
            return outputs, non_trainable_variables, scope.losses[:]
        return outputs, non_trainable_variables

    def compute_output_spec(self, *args, **kwargs):
        if utils.is_default(self.compute_output_shape):
            return super().compute_output_spec(*args, **kwargs)
        else:
            
            call_spec = CallSpec(self._call_signature, args, kwargs)
            shapes_dict = get_shapes_dict(call_spec)
            shapes_dict = update_shapes_dict_for_target_fn(
                self.compute_output_shape,
                shapes_dict=shapes_dict,
                call_spec=call_spec,
                class_name=self.__class__.__name__,
            )
            output_shape = self.compute_output_shape(**shapes_dict)

            if (
                isinstance(output_shape, list)
                and output_shape
                and isinstance(output_shape[0], (int, type(None)))
            ):
                output_shape = tuple(output_shape)
            if not isinstance(output_shape, (list, tuple, dict)):
                try:
                    output_shape = tuple(output_shape)
                except:
                    raise ValueError(
                        "Method `compute_output_shape()` of layer "
                        f"{self.__class__.__name__} is returning "
                        "a type that cannot be interpreted as a shape. "
                        "It should return a shape tuple. "
                        f"Received: {output_shape}"
                    )
            if (
                isinstance(output_shape, tuple)
                and output_shape
                and isinstance(output_shape[0], (int, type(None)))
            ):
                return KerasTensor(output_shape, dtype=self.compute_dtype)
            
            
            return map_shape_structure(
                lambda s: KerasTensor(s, dtype=self.compute_dtype), output_shape
            )

    @utils.default
    def compute_output_shape(self, *args, **kwargs):
        raise NotImplementedError(
            f"Layer {self.__class__.__name__} should implement "
            "`def compute_output_shape(self, input_shape)`."
        )

    def add_loss(self, loss):
        
        
        losses = tree.flatten(loss)
        for x in losses:
            if not backend.is_tensor(x):
                raise ValueError(
                    "`add_loss()` can only be called from inside `build()` or "
                    f"`call()`, on a tensor input. Received invalid value: {x}"
                )
        if backend.in_stateless_scope():
            scope = backend.get_stateless_scope()
            if scope.collect_losses:
                for x in losses:
                    scope.add_loss(loss)
                    self._loss_ids.add(id(loss))
        else:
            self._losses.extend(losses)

    def _get_own_losses(self):
        if backend.in_stateless_scope():
            losses = []
            scope = backend.get_stateless_scope()
            for loss in scope.losses:
                if id(loss) in self._loss_ids:
                    losses.append(loss)
            return losses
        else:
            return self._losses[:]

    @property
    def losses(self):
        
        losses = self._get_own_losses()
        for layer in self._flatten_layers(include_self=False):
            losses.extend(layer._get_own_losses())
        weight_regularization_losses = []
        for v in self.trainable_weights:
            if backend.in_stateless_scope():
                v = backend.get_stateless_scope().get_current_value(v)
            regularizer = getattr(v, "regularizer", None)
            if regularizer:
                weight_regularization_losses.append(regularizer(v))
        losses.extend(weight_regularization_losses)
        return losses

    def _clear_losses(self):
        if backend.in_stateless_scope():
            scope = backend.get_stateless_scope()
            if scope.collect_losses:
                for x in scope.losses:
                    if id(x) in self._loss_ids:
                        scope.losses.remove(x)
        self._losses.clear()
        self._loss_ids.clear()
        for layer in self._layers:
            layer._clear_losses()

    def save_own_variables(self, store):
        
        all_vars = self._trainable_variables + self._non_trainable_variables
        for i, v in enumerate(all_vars):
            store[f"{i}"] = v.numpy()

    def load_own_variables(self, store):
        
        all_vars = self._trainable_variables + self._non_trainable_variables
        if len(store.keys()) != len(all_vars):
            if len(all_vars) == 0 and not self.built:
                raise ValueError(
                    f"Layer '{self.name}' was never built "
                    "and thus it doesn't have any variables. "
                    f"However the weights file lists {len(store.keys())} "
                    "variables for this layer.\n"
                    "In most cases, this error indicates that either:\n\n"
                    "1. The layer is owned by a parent layer that "
                    "implements a `build()` method, but calling the "
                    "parent's `build()` method did NOT create the state of "
                    f"the child layer '{self.name}'. A `build()` method "
                    "must create ALL state for the layer, including "
                    "the state of any children layers.\n\n"
                    "2. You need to implement "
                    "the `def build_from_config(self, config)` method "
                    f"on layer '{self.name}', to specify how to rebuild "
                    "it during loading. "
                    "In this case, you might also want to implement the "
                    "method that generates the build config at saving time, "
                    "`def get_build_config(self)`. "
                    "The method `build_from_config()` is meant "
                    "to create the state "
                    "of the layer (i.e. its variables) upon deserialization.",
                )
            raise ValueError(
                f"Layer '{self.name}' expected {len(all_vars)} variables, "
                "but received "
                f"{len(store.keys())} variables during loading. "
                f"Expected: {[v.name for v in all_vars]}"
            )
        for i, v in enumerate(all_vars):
            v.assign(store[f"{i}"])

    def _track_variable(self, variable):
        if variable.trainable:
            self._tracker.add_to_store("trainable_variables", variable)
        else:
            self._tracker.add_to_store("non_trainable_variables", variable)

    def add_metric(self):
        
        raise NotImplementedError

    def count_params(self):
        
        if not self.built:
            raise ValueError(
                "You tried to call `count_params` "
                f"on layer '{self.name}', "
                "but the layer isn't built. "
                "You can build it manually via: "
                f"`layer.build(input_shape)`."
            )
        return summary_utils.count_params(self.weights)

    def _maybe_build(self, call_spec):
        if self.built:
            return

        shapes_dict = get_shapes_dict(call_spec)
        first_shape = next(iter(shapes_dict.values()), None)

        
        if not utils.is_default(self.build):
            shapes_dict = update_shapes_dict_for_target_fn(
                self.build,
                shapes_dict=shapes_dict,
                call_spec=call_spec,
                class_name=self.__class__.__name__,
            )
            self.build(**shapes_dict)
            
            
            self._assert_input_compatibility(call_spec.first_arg)
            return

        
        if might_have_unbuilt_state(self):
            try:
                backend.compute_output_spec(
                    self.call, **call_spec.arguments_dict
                )
            except Exception as e:
                if call_spec.eager:
                    
                    return
                warnings.warn(
                    f"Layer '{self.name}' looks like it has unbuilt state, but "
                    "Keras is not able to trace the layer `call()` in order to "
                    "build it automatically. Possible causes:\n"
                    "1. The `call()` method of your layer may be crashing. Try "
                    "to `__call__()` the layer eagerly on some test input "
                    "first to see if it works. "
                    "E.g. `x = np.random.random((3, 4)); y = layer(x)`\n"
                    "2. If the `call()` method is correct, then you may need "
                    "to implement the `def build(self, input_shape)` method on "
                    "your layer. It should create all variables used by the "
                    "layer (e.g. by calling `layer.build()` on all its "
                    "children layers).\n"
                    f"Exception encoutered: ''{e}''"
                )
        self.build(first_shape)

    def _build_by_run_for_single_pos_arg(self, input_shape):
        
        input_tensors = map_shape_structure(
            lambda s: backend.KerasTensor(s), input_shape
        )
        try:
            backend.compute_output_spec(self.call, input_tensors)
            return True
        except:
            return False

    def _build_by_run_for_kwargs(self, shapes_dict):
        
        if all(is_shape_tuple(s) for s in shapes_dict.values()):
            
            input_tensors = {
                
                utils.removesuffix(k, "_shape"): backend.KerasTensor(shape)
                for k, shape in shapes_dict.items()
            }
            try:
                backend.compute_output_spec(self.call, **input_tensors)
                return True
            except:
                return False
        else:
            
            return False

    def __repr__(self):
        return (
            f"<{self.__class__.__name__} "
            f"name={self.name}, built={self.built}>"
        )

    def __str__(self):
        return (
            f"<{self.__class__.__name__} "
            f"name={self.name}, built={self.built}>"
        )

    def __setattr__(self, name, value):
        
        name, value = self._setattr_hook(name, value)
        if hasattr(self, "_tracker"):
            value = self._tracker.track(value)
        elif name != "_tracker":
            self._initializer_tracker()
        return super().__setattr__(name, value)

    def _check_super_called(self):
        if getattr(self, "_lock", True):
            raise RuntimeError(
                f"In layer '{self.__class__.__name__}', you forgot to call "
                "`super().__init__()` as the first statement "
                "in the `__init__()` method. Go add it!"
            )

    def _assert_input_compatibility(self, arg_0):
        if self.input_spec:
            input_spec.assert_input_compatibility(
                self.input_spec, arg_0, layer_name=self.name
            )

    def _get_call_context(self):
        
        layer_call_ctx = global_state.get_global_attribute("current_call_ctx")
        if layer_call_ctx is None:
            
            layer_call_ctx = CallContext(entry_layer=self)
            global_state.set_global_attribute(
                "current_call_ctx", layer_call_ctx
            )
            self._clear_losses()
        return layer_call_ctx

    def _maybe_reset_call_context(self):
        layer_call_ctx = global_state.get_global_attribute("current_call_ctx")
        if layer_call_ctx is None or layer_call_ctx.entry_layer == self:
            global_state.set_global_attribute("current_call_ctx", None)

    def _flatten_layers(self, include_self=True, recursive=True):
        layers = []
        if include_self:
            layers.append(self)
        seen_object_ids = set()
        deque = collections.deque(self._layers)
        while deque:
            layer = deque.popleft()
            if id(layer) in seen_object_ids:
                continue
            seen_object_ids.add(id(layer))
            layers.append(layer)
            
            if recursive:
                deque.extendleft(layer._layers)
        return layers

    def _set_mask_metadata(self, inputs, outputs, previous_mask):
        flat_outputs = tree.flatten(outputs)

        mask_already_computed = all(
            getattr(x, "_keras_mask", None) is not None for x in flat_outputs
        )
        if mask_already_computed:
            return

        output_masks = self.compute_mask(inputs, previous_mask)
        if output_masks is None:
            return

        flat_masks = tree.flatten(output_masks)
        for tensor, mask in zip(flat_outputs, flat_masks):
            if getattr(tensor, "_keras_mask", None) is None:
                try:
                    
                    if backend.backend() == "numpy":
                        warnings.warn(
                            "The NumPy backend does not support masking at this"
                            "time. Masks will be ignored."
                        )
                    tensor._keras_mask = mask
                except AttributeError:
                    
                    pass

    @python_utils.default
    def get_config(self):
        self._check_super_called()
        base_config = super().get_config()
        config = {
            "trainable": self.trainable,
            "dtype": self.dtype_policy.name,
        }
        return {**base_config, **config}


def is_backend_tensor_or_symbolic(x):
    return backend.is_tensor(x) or isinstance(x, backend.KerasTensor)


class CallSpec:
    def __init__(self, signature, args, kwargs):
        
        
        
        
        
        if "training" in kwargs and "training" not in signature.parameters:
            kwargs.pop("training")
            bound_args = signature.bind(*args, **kwargs)
        else:
            bound_args = signature.bind(*args, **kwargs)
        self.user_arguments_dict = {
            k: v for k, v in bound_args.arguments.items()
        }
        bound_args.apply_defaults()
        arg_dict = {}
        arg_names = []
        tensor_arg_dict = {}
        tensor_args = []
        tensor_arg_names = []
        nested_tensor_arg_names = []
        for name, value in bound_args.arguments.items():
            arg_dict[name] = value
            arg_names.append(name)
            if is_backend_tensor_or_symbolic(value):
                tensor_args.append(value)
                tensor_arg_names.append(name)
                tensor_arg_dict[name] = value
            elif tree.is_nested(value):
                flat_values = tree.flatten(value)
                if all(is_backend_tensor_or_symbolic(x) for x in flat_values):
                    tensor_args.append(value)
                    tensor_arg_names.append(name)
                    tensor_arg_dict[name] = value
                    nested_tensor_arg_names.append(name)
                elif any(is_backend_tensor_or_symbolic(x) for x in flat_values):
                    raise ValueError(
                        "In a nested call() argument, "
                        "you cannot mix tensors and non-tensors. "
                        "Received invalid mixed argument: "
                        f"{name}={value}"
                    )
        self.arguments_dict = arg_dict
        self.argument_names = arg_names
        self.tensor_arguments_dict = tensor_arg_dict
        self.tensor_arguments_names = tensor_arg_names
        self.nested_tensor_argument_names = nested_tensor_arg_names
        self.first_arg = arg_dict[arg_names[0]]
        if all(
            backend.is_tensor(x) for x in self.tensor_arguments_dict.values()
        ):
            self.eager = True
        else:
            self.eager = False


def get_arguments_dict(fn, args, kwargs):
    
    sig = inspect.signature(fn)
    bound_args = sig.bind(*args, **kwargs)
    arg_dict = {}
    for name, value in bound_args.arguments.items():
        arg_dict[name] = value
    return arg_dict


def get_shapes_dict(call_spec):
    
    shapes_dict = {}
    for k, v in call_spec.tensor_arguments_dict.items():
        if k == "mask" or k.endswith("_mask"):
            
            continue
        if k == "kwargs" or k == "args":
            
            continue
        if k in call_spec.nested_tensor_argument_names:
            shapes_dict[f"{k}_shape"] = tree.map_structure(
                lambda x: backend.standardize_shape(x.shape), v
            )
        else:
            shapes_dict[f"{k}_shape"] = backend.standardize_shape(v.shape)
    return shapes_dict


def update_shapes_dict_for_target_fn(
    target_fn,
    shapes_dict,
    call_spec,
    class_name,
):
    
    if utils.is_default(target_fn):
        return None
    sig = inspect.signature(target_fn)
    expected_names = []
    for name, param in sig.parameters.items():
        if param.kind in (
            param.POSITIONAL_OR_KEYWORD,
            param.POSITIONAL_ONLY,
            param.KEYWORD_ONLY,
        ):
            expected_names.append(name)

    
    if len(expected_names) == 1:
        key = expected_names[0]
        values = tuple(shapes_dict.values())
        if values:
            input_shape = values[0]
        else:
            input_shape = None
        return {key: input_shape}

    
    kwargs = {}
    for name in expected_names:
        method_name = target_fn.__name__
        error_preamble = (
            f"For a `{method_name}()` method with more than one argument, all "
            "arguments should have a `_shape` suffix and match an argument "
            f"from `call()`. E.g. `{method_name}(self, foo_shape, bar_shape)` "
        )
        if not name.endswith("_shape"):
            raise ValueError(
                f"{error_preamble} For layer '{class_name}', "
                f"Received `{method_name}()` argument "
                f"`{name}`, which does not end in `_shape`."
            )
        expected_call_arg = utils.removesuffix(name, "_shape")
        if expected_call_arg not in call_spec.arguments_dict:
            raise ValueError(
                f"{error_preamble} For layer '{class_name}', "
                f"received `{method_name}()` argument "
                f"`{name}`, but `call()` does not have argument "
                f"`{expected_call_arg}`."
            )
        if name in shapes_dict:
            kwargs[name] = shapes_dict[name]

    return kwargs


class CallContext:
    def __init__(self, entry_layer):
        self.entry_layer = entry_layer
        self.training = None


def is_shape_tuple(s):
    return isinstance(s, (list, tuple)) and all(
        d is None or isinstance(d, int) for d in s
    )


def might_have_unbuilt_state(layer):
    return any(not lr.built for lr in layer._layers)