Code update

PiperOrigin-RevId: 561839173

swirl_dynamics/projects/ergodic/utils.py

@@ -304,3 +304,19 @@ def sample_uniform_spherical_shell(
   vec /= vec_norm
 
   return vec * sampling_radius.reshape(broadcasting_shape)
+
+
+def linear_scale_dissipative_target(inputs: Array, scale: float = 1.0):
+  """Function to rescale the random input to bias the model.
+
+  Arguments:
+    inputs: Input point in state space.
+    scale: Real number in [0, 1] that scales down input.
+
+  Returns:
+    The rescaled input, scale*input.
+
+  This function is implemented to follow the implementation of the Markov Neural
+  Operator (MNO).
+  """
+  return scale * inputs