from keras.src.api_export import keras_export

from keras.src.backend.common import global_state

@keras_export("keras.StatelessScope")

class StatelessScope:

"""Scope to prevent any update to Keras Variables.

The values of variables to be used inside the scope

should be passed via the `state_mapping` argument, a

list of tuples `(k, v)` where `k` is a `KerasVariable`

and `v` is the intended value for this variable

(a backend tensor).

Updated values can be collected on scope exit via

`value = scope.get_current_value(variable)`. No updates

will be applied in-place to any variables for the duration

of the scope.

Example:

```python

state_mapping = [(k, ops.ones(k.shape, k.dtype)) for k in model.weights]

with keras.StatelessScope(state_mapping) as scope:

outputs = model.some_function(inputs)

# All model variables remain unchanged. Their new values can be

# collected via:

for k in model.weights:

new_value = scope.get_current_value(k)

print(f"New value for {k}: {new_value})

```

"""

def __init__(

self,

state_mapping=None,

collect_losses=False,

initialize_variables=True,

):

from keras.src import backend

from keras.src.backend.common.variables import KerasVariable

self.collect_losses = collect_losses

self.initialize_variables = initialize_variables

self.losses = []

self.state_mapping = {}

state_mapping = state_mapping or {}

for k, v in state_mapping:

if not isinstance(k, KerasVariable):

raise ValueError(

"Invalid reference variable in StatelessScope: "

"all keys in argument `mapping` must be KerasVariable "

f"instances. Received instead: {k}"

)

if isinstance(v, KerasVariable):

v = backend.cast(v.value, dtype=k.dtype)

else:

v = backend.convert_to_tensor(v, dtype=k.dtype)

if k.shape != v.shape:

raise ValueError(

"Invalid variable value in StatelessScope: "

"all values in argument `mapping` must be tensors with "

"a shape that matches the corresponding variable shape. "

f"For variable {k}, received invalid value {v} with shape "

f"{v.shape}."

)

self.state_mapping[id(k)] = v

def __enter__(self):

self.original_scope = get_stateless_scope()

global_state.set_global_attribute("stateless_scope", self)

return self

def add_loss(self, loss):

self.losses.append(loss)

def add_update(self, update):

variable, value = update

self.state_mapping[id(variable)] = value

def get_current_value(self, variable):

return self.state_mapping.get(id(variable), None)

def __exit__(self, *args, **kwargs):

global_state.set_global_attribute(

"stateless_scope", self.original_scope

)

if self.original_scope is None and self.initialize_variables:

# We're back in eager scope;

# if any variables were created within the stateless

# scope, we initialize them here.

from keras.src.backend.common.variables import (

initialize_all_variables,

)

initialize_all_variables()

def in_stateless_scope():

return global_state.get_global_attribute("stateless_scope") is not None

def get_stateless_scope():

return global_state.get_global_attribute("stateless_scope")